Creating Invariant Tests for an AMM Smart Contract

Invariant testing in smart contracts is about repeatedly trying out different combinations of actions to make sure the contract always follows its core rules, even in unusual or unexpected situations.

Source: https://www.mathsisfun.com/definitions/invariant.html

What is Invariant Testing?

Invariant testing in smart contracts is about repeatedly trying out different combinations of actions to make sure the contract always follows its core rules, even in unusual or unexpected situations.

It’s like playing a complex card game where you shuffle and deal the cards randomly and then play them in random order. Each card represents a command or action in a smart contract. After playing each card (or performing each action), you check some basic rules to make sure the game still makes sense.

In Foundry’s context, consider each “run” of a test as a separate session of a card game, with each session comprising multiple rounds (or “calls”).

During each session, the game is played multiple times, and different strategies (test scenarios) are employed in each round. The aim is to explore various outcomes and ensure that the rules (invariants) of the game are always upheld, regardless of the strategies used.

Difference between Invariant and Fuzz Testing

Invariant testing is about ensuring internal consistency and adherence to predefined logical conditions, while fuzz testing is about assessing how software handles unexpected and potentially invalid external inputs, with a strong emphasis on identifying security issues and vulnerabilities.

Let’s go through some points:

Objective and Focus:

• Invariant Testing: The primary objective is to ensure that certain conditions, or invariants, always hold true throughout the execution of a program, regardless of the program’s state or the inputs it receives. 

• Fuzz Testing: Fuzz testing aims to uncover vulnerabilities and ensure robustness by feeding invalid, unexpected, or random data as inputs. 

Methodology:

• Invariant Testing: Involves defining specific invariants in the software and then repeatedly checking these invariants at runtime to ensure they are always true.

• Fuzz Testing: Employs the generation of random or pseudo-random data (the “fuzz”) which is then used as input to the system. 

Nature of Testing:

• Invariant Testing: It’s deterministic in nature, as it relies on specific, predefined conditions that must always be met.

• Fuzz Testing: Non-deterministic and more exploratory, since it involves generating a wide range of unpredictable inputs to test how the system responds to them, particularly focusing on edge cases and unexpected scenarios.

How do we define invariants?

When identifying invariants, it’s essential to consider the contract’s fundamental purpose, how its functions interact with each other, and the broader financial or operational logic it embodies. 

Each function, variable, and event can give insights into what should remain constant, what conditions must be met, and what the contract aims to achieve or maintain at all times. 

This understanding then guides the creation of invariants for testing and ensuring the contract’s reliability and correctness.

Learning invariant testing might get complicated and defining those invariants from an existing code is tricky, so we are going to implement a simple AMM-based contract and, before testing it, we are going to see which are its invariants and how to define them to prepare for testing them.

Here is our SimpleAMM contract that we will use for the sole purpose of learning about Invariant Testing:

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract SimpleAMM {
    uint256 public reserveTokenA;
    uint256 public reserveTokenB;
    uint256 public constantProduct;

    event LiquidityAdded(address indexed user, uint256 tokenAAmount, uint256 tokenBAmount);
    event LiquidityRemoved(address indexed user, uint256 tokenAAmount, uint256 tokenBAmount);
    event TokensSwapped(address indexed user, uint256 tokenASwapped, uint256 tokenBReceived);

    function addLiquidity(uint256 tokenAAmount, uint256 tokenBAmount) external {
        require(tokenAAmount > 0 && tokenBAmount > 0, "Cannot add zero liquidity");
        reserveTokenA += tokenAAmount;
        reserveTokenB += tokenBAmount;
        updateConstantProduct();
        emit LiquidityAdded(msg.sender, tokenAAmount, tokenBAmount);
    }

    function removeLiquidity(uint256 tokenAAmount, uint256 tokenBAmount) external {
        require(reserveTokenA >= tokenAAmount && reserveTokenB >= tokenBAmount, "Insufficient liquidity");
        reserveTokenA -= tokenAAmount;
        reserveTokenB -= tokenBAmount;
        updateConstantProduct();
        emit LiquidityRemoved(msg.sender, tokenAAmount, tokenBAmount);
    }

    function swapTokenAForTokenB(uint256 tokenAAmount) external {
        uint256 tokenBAmount = getSwapAmount(tokenAAmount);
        require(reserveTokenB >= tokenBAmount, "Insufficient token B reserve");
        reserveTokenA += tokenAAmount;
        reserveTokenB -= tokenBAmount;
        updateConstantProduct();
        emit TokensSwapped(msg.sender, tokenAAmount, tokenBAmount);
    }

    function getSwapAmount(uint256 tokenAAmount) public view returns (uint256) {
        uint256 newReserveA = reserveTokenA + tokenAAmount;
        require(newReserveA > 0, "Invalid swap amount");
        uint256 newReserveB = constantProduct / newReserveA;
        return reserveTokenB - newReserveB;
    }

    function _updateConstantProduct() private {
        constantProduct = reserveTokenA * reserveTokenB;
    }
}

Automated Market Makers are a crucial component of DeFi, and while each AMM can have its unique mechanisms and features, there are several general invariants that typically apply across different AMM implementations. 

I won’t elaborate on them, but I consider it relevant to mention some of them in this short list:

1. Reserve Ratio or Pricing Formula Compliance 

2. Non-Negative Reserves

3. Liquidity Pool Solvency

4. Liquidity Provider (LP) Token Consistency

5. Fair Exchange Rates

6. Slippage and Price Impact Limits

7. Fee Accrual and Distribution

8. Impermanent Loss Constraints

9. Arbitrage Consistency

10. Token Conservation

Now it’s time to start defining the invariants for our SimpleAMM contract, so, let’s see where and how can we extract those invariants.

SimpleAMM contract’s invariants

1. Constant Product Invariant 
   After any operation (adding/removing liquidity, swapping tokens), the product of reserveTokenA and reserveTokenB should remain equal to constantProduct.

• Purpose: To ensure that the AMM maintains its defining characteristic of a constant product market maker.

• Implementation with Foundry:

function invariant_ConstantProduct() public {
    uint256 reserveTokenA = simpleAMM.reserveTokenA();
    uint256 reserveTokenB = simpleAMM.reserveTokenB();
    uint256 constantProduct = simpleAMM.constantProduct();
    assertEq(
        reserveTokenA * reserveTokenB,
        constantProduct,
        "Constant product invariant violated"
    );
}

2. Reserve Non-Negativity
The reserves reserveTokenA and reserveTokenB should never become negative. This ensures the contract does not end up in a state where it owes more tokens than it holds.

• Purpose: To verify that these operations do not result in negative reserves, which would be physically and financially infeasible.

• Implementation with Foundry:

// Invariant: Reserve Non-Negativity
function invariant_ReserveNonNegativity() public {
    assertGe(simpleAMM.reserveTokenA(), 0, "reserveTokenA is negative");
    assertGe(simpleAMM.reserveTokenB(), 0, "reserveTokenB is negative");
}

3. Swap Rate Fairness
The rate at which tokens are swapped (getSwapAmount) should always be fair according to the current reserves. This means after a swap, the product of the reserves (adjusted for any fees, if applicable) should stay the same or very close to the constantProduct.

• Purpose: Ensures that the rate at which tokens are swapped is fair and consistent with the constant product formula.

• Implementation with Foundry:

// Invariant: Swap Rate Fairness
function invariant_SwapRateFairness() public {
    uint256 reserveTokenA = simpleAMM.reserveTokenA();
    uint256 reserveTokenB = simpleAMM.reserveTokenB();
    uint256 constantProduct = simpleAMM.constantProduct();

    // Allow for a small tolerance due to integer division rounding
    uint256 tolerance = 1; // Define an appropriate tolerance

    uint256 product = reserveTokenA * reserveTokenB;
    bool isProductCloseToConstant = (product >=
        constantProduct - tolerance) &&
        (product <= constantProduct + tolerance);

    assertTrue(
        isProductCloseToConstant,
        "Swap rate fairness violated: Product of reserves not close to constant product"
    );
}

---

At this stage, I have written a Handler contract to be able to add restrictions to my tests.

Let’s, first of all, go through what is a Handler and why I needed to use it.

Using a Handler Contract

A handler contract acts as an intermediary layer between my tests and the SimpleAMM contract. It's essentially another smart contract that wraps the functions of SimpleAMM.

Why Use a Handler?

• Modularity: 
  A handler can encapsulate specific behaviors or sequences of interactions. This is particularly useful when you have complex interactions or need to maintain certain states between calls.

• Reusability: 
  You can reuse the same handler logic across different tests or even different test suites. This can be more efficient if you have a set of standard operations you perform in multiple tests.

• Complex Interactions: 
  If your testing involves complex sequences or setups (like preparing a certain state before calling a function), a handler can abstract away these complexities from the test itself.

• Boundaries:
  Often in order to properly test your invariant and avoid false failures you will have to specify some maximum and/or minimum integers to restrict the possibilities for the fuzzer.

And here is my SimpleAMMHandler contract:

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "../../src/SimpleAMM.sol";
import {CommonBase} from "forge-std/Base.sol";
import {StdCheats} from "forge-std/StdCheats.sol";
import {StdUtils} from "forge-std/StdUtils.sol";

contract SimpleAMMHandler is CommonBase, StdCheats, StdUtils {
    SimpleAMM public simpleAMM;

    constructor(SimpleAMM _simpleAMM) {
        simpleAMM = _simpleAMM;
        simpleAMM.addLiquidity(1000, 1000);
    }

    function addLiquidity(uint256 tokenAAmount, uint256 tokenBAmount) external {
        tokenAAmount = bound(tokenAAmount, 1, 10000);
        tokenBAmount = bound(tokenBAmount, 1, 10000);

        simpleAMM.addLiquidity(tokenAAmount, tokenBAmount);
    }

    function removeLiquidity(
        uint256 tokenAAmount,
        uint256 tokenBAmount
    ) external {
        tokenAAmount = bound(tokenAAmount, 1, 10000);
        tokenBAmount = bound(tokenBAmount, 1, 10000);

        simpleAMM.removeLiquidity(tokenAAmount, tokenBAmount);
    }

    function swapTokenAForTokenB(uint256 tokenAAmount) external {
        tokenAAmount = bound(tokenAAmount, 1, 10000);

        simpleAMM.swapTokenAForTokenB(tokenAAmount);
    }
}

---

4. Liquidity Addition and Removal Consistency
Adding liquidity and then removing the same amount should leave the reserves unchanged (ignoring any transaction fees or other external factors that might affect the reserves).

• Purpose: To confirm that adding and then removing the same amount of liquidity leaves the reserves unchanged, barring transaction fees or other factors.

• Implementation with Foundry:

// Invariant: Liquidity Addition and Removal Consistency
function invariant_LiquidityAdditionRemovalConsistency() public {
    uint256 initialReserveA = simpleAMM.reserveTokenA();
    uint256 initialReserveB = simpleAMM.reserveTokenB();

    // Perform a test liquidity addition and removal
    uint256 testAmountA = 500; // Choose appropriate test amounts
    uint256 testAmountB = 500;
    simpleAMMHandler.addLiquidity(testAmountA, testAmountB);
    simpleAMMHandler.removeLiquidity(testAmountA, testAmountB);

    uint256 finalReserveA = simpleAMM.reserveTokenA();
    uint256 finalReserveB = simpleAMM.reserveTokenB();

    // Assert the final reserves are equal to the initial reserves
    assertEq(
        finalReserveA,
        initialReserveA,
        "Reserve A should remain unchanged after add/remove liquidity"
    );
    assertEq(
        finalReserveB,
        initialReserveB,
        "Reserve B should remain unchanged after add/remove liquidity"
    );
}

5. No Token Creation or Destruction
The contract should not create or destroy tokens. The sum of the tokens in the reserves plus the tokens held by users should always equal the total supply of the tokens.

• Purpose: To ensure the contract does not inadvertently create or destroy tokens, which would violate the conservation of token supply.

• Implementation with Foundry:

// Invariant: No Token Creation or Destruction
function invariant_TokenConservation() public {
    // Keep in mind this is a simplified AMM with only one known pair 
    IERC20 tokenA = IERC20(tokenAAddress);
    IERC20 tokenB = IERC20(tokenBAddress);

    uint256 totalSupplyA = tokenA.totalSupply();
    uint256 totalSupplyB = tokenB.totalSupply();

    uint256 reserveA = simpleAMM.reserveTokenA();
    uint256 reserveB = simpleAMM.reserveTokenB();

    // Calculating total tokens held by users (excluding reserves)
    // This could be handled in multiple ways and since it's not relevant for
    // the sake of our invariant test exercise, assume it does what it says
    uint256 userHeldA = calculateUserHeldTokens(tokenA);
    uint256 userHeldB = calculateUserHeldTokens(tokenB);

    // Asserting that total supply equals sum of reserves and user-held tokens
    assertEq(totalSupplyA, reserveA + userHeldA, "Token A conservation violated");
    assertEq(totalSupplyB, reserveB + userHeldB, "Token B conservation violated");
}

6. Positive Liquidity
The contract should not allow the addition of zero liquidity. Both tokenAAmount and tokenBAmount in addLiquidity must be greater than zero.

• Purpose: Ensures that liquidity addition is meaningful (non-zero) and aligns with the economic logic of the AMM.

• Implementation with Foundry:

// Invariant: Positive Liquidity
function invariant_PositiveLiquidity() public {
    uint256 reserveA = simpleAMM.reserveTokenA();
    uint256 reserveB = simpleAMM.reserveTokenB();

    // Check that both reserves are greater than zero
    assertTrue(reserveA > 0, "Liquidity for token A is not positive");
    assertTrue(reserveB > 0, "Liquidity for token B is not positive");
}

7. Swap Amount Validation
The amount returned by getSwapAmount must be less than or equal to the current reserveTokenB when swapping Token A for Token B.

• Purpose: To confirm that the calculated swap amount is always within the available reserves.

• Implementation with Foundry:

// Invariant: Swap Amount Validation
function invariant_SwapAmountValidation() public {
    uint256 reserveTokenB = simpleAMM.reserveTokenB();
    // Generate a random amount of Token A
    uint256 randomAmountTokenA = 1234; // Let's assume this hardcoded number is random one

    uint256 swapAmountB = simpleAMM.getSwapAmount(randomAmountTokenA);

    // Check that both reserves are greater than zero
    assertTrue(
        swapAmountB <= reserveTokenB,
        "Swap amount exceeds reserveTokenB"
    );
}

8. Solvency
After any swap, the contract must remain solvent, meaning the reserves must be sufficient to satisfy the constant product formula.

• Purpose: To ensure that after any operation, the contract remains solvent and can fulfill its obligations according to the constant product rule.

• Implementation with Foundry:

// Invariant: Solvency
function invariant_Solvency() public {
    uint256 reserveTokenA = simpleAMM.reserveTokenA();
    uint256 reserveTokenB = simpleAMM.reserveTokenB();
    uint256 minimumReserveThreshold = 500; // Randomly chosen for the sake of example

    // Ensure reserves are above a minimum threshold
    assertTrue(
        reserveTokenA > minimumReserveThreshold,
        "Insufficient ReserveTokenA for solvency"
    );
    assertTrue(
        reserveTokenB > minimumReserveThreshold,
        "Insufficient ReserveTokenB for solvency"
    );
}

---

WATCH OUT! 
WE HAVE FOUND A VULNERABILITY ON OUR CONTRACT!!!

Yes, that’s correct, there was nothing protecting our SimpleAMM contract of assuring solvency and this invariant test has just pointed that out.

In order to mitigate this we will have to add some logic to our removeLiquidity and swapTokenAForTokenB functions, such as:

function swapTokenAForTokenB(uint256 tokenAAmount) external {
    // ... existing logic ...

    // Check final reserves after swap
    require(reserveTokenA() >= 500, "Reserve Token A below minimum threshold");
    require(reserveTokenB() >= 500, "Reserve Token B below minimum threshold");

    // ... rest of the function ...
}

function removeLiquidity(uint256 tokenAAmount, uint256 tokenBAmount) external {
    // ... existing logic ...

    // Check final reserves after liquidity removal
    require(reserveTokenA() >= 500, "Reserve Token A below minimum threshold");
    require(reserveTokenB() >= 500, "Reserve Token B below minimum threshold");

    // ... rest of the function ...
}

---

9. Event Emission Consistency 
Ensure that every state-changing function emits the correct event with accurate parameters reflecting the state change.

• Purpose: To guarantee that the contract accurately reflects state changes through events, which is crucial for transparency and tracking.

• Implementation in Foundry? In this case, from my point of view, this is an invariant that will need to be done manually, and yet it is as important as the rest.

10. User Balance Consistency
The balance of a user’s tokens should reflect the sum of their provided liquidity plus any tokens not deposited in the contract. This is more of an external invariant but still crucial for end-to-end testing.

• Purpose: To ensure that the contract’s operations correctly reflect on the overall token balances of users, maintaining the integrity of user holdings.

• Implementation in Foundry: [Keep reading, time to learn something new]

---

SPOILER ALERT
A new cool fuzzing concept is coming…

We can agree that it would be complicated or maybe not even possible to simulate such a scenario to control the flow of added, removed, and swapped tokens.

At this stage is time to bring in and start using the so-called ghost variables. Let’s see what they can do for us and how are we going to use it to test the user balance consistency invariant.

Ghost variables are essentially state variables that are used only for testing purposes. They’re not part of the actual contract logic but are added to the contract during testing to keep track of certain states or values that are otherwise not explicitly stored in the contract.

Where do we add them?

Definitely not in the production version of our smart contract. So, the next best thing is the Handler we created earlier, isn’t it?

We can use ghost variables in our SimpleAMMHandler contract to track a state that is not otherwise exposed by the contract under test.

Below is how our contract looks after adding the ghost variables. Keep an eye on ghost_expectedBalanceTokenA and ghost_expectedBalanceTokenB.

contract SimpleAMMHandler is CommonBase, StdCheats, StdUtils {
    SimpleAMM public simpleAMM;

    // Ghost variables to track expected user balances
    mapping(address => uint256) public ghost_expectedBalanceTokenA;
    mapping(address => uint256) public ghost_expectedBalanceTokenB;

    constructor(SimpleAMM _simpleAMM) {
        simpleAMM = _simpleAMM;
        simpleAMM.addLiquidity(1000, 1000);

        // Let's give our msg.sender some credit
        ghost_expectedBalanceTokenA[msg.sender] += 5000;
        ghost_expectedBalanceTokenB[msg.sender] += 5000;
    }

    function addLiquidity(uint256 tokenAAmount, uint256 tokenBAmount) external {
        tokenAAmount = bound(tokenAAmount, 1, 10000);
        tokenBAmount = bound(tokenBAmount, 1, 10000);

        // Adjust expected balances
        ghost_expectedBalanceTokenA[msg.sender] -= tokenAAmount;
        ghost_expectedBalanceTokenB[msg.sender] -= tokenBAmount;

        simpleAMM.addLiquidity(tokenAAmount, tokenBAmount);
    }

    function removeLiquidity(
        uint256 tokenAAmount,
        uint256 tokenBAmount
    ) external {
        tokenAAmount = bound(tokenAAmount, 1, 10000);
        tokenBAmount = bound(tokenBAmount, 1, 10000);

        // Adjust expected balances
        ghost_expectedBalanceTokenA[msg.sender] += tokenAAmount;
        ghost_expectedBalanceTokenB[msg.sender] += tokenBAmount;

        simpleAMM.removeLiquidity(tokenAAmount, tokenBAmount);
    }

    function swapTokenAForTokenB(uint256 tokenAAmount) external {
        uint256 tokenBAmount;
        tokenAAmount = bound(tokenAAmount, 1, 10000);

        // Adjust expected balances
        ghost_expectedBalanceTokenA[msg.sender] -= tokenAAmount;
        ghost_expectedBalanceTokenB[msg.sender] += tokenBAmount;

        simpleAMM.swapTokenAForTokenB(tokenAAmount);
    }
}

Also, we will need a way to retrieve the list of user addresses, so in this case, we will add the state variable address[] private userAddresses and the functions below at the end of our handler:

// Function to retrieve user addresses
function getUserAddresses() external view returns (address[] memory) {
    return userAddresses;
}

// Utility function to add a new user address to the list
function addUserAddress(address user) internal {
    if (!addressExists(user)) {
        userAddresses.push(user);
    }
}

// Check if an address is already in the list
function addressExists(address user) internal view returns (bool) {
    for (uint i = 0; i < userAddresses.length; i++) {
        if (userAddresses[i] == user) {
            return true;
        }
    }
    return false;
}

---

Let’s use now the ghost variables in our tests to track user balances and liquidity provided over time. The idea here is to assert that the expected token balances (tracked by the ghost variables in the handler) always match the actual token balances of users.

// Invariant: User Balance Consistency
function invariant_UserBalanceConsistency() public {
    address[] memory userAddresses = simpleAMMHandler.getUserAddresses();

    for (uint i = 0; i < userAddresses.length; i++) {
        address user = userAddresses[i];

        uint256 expectedBalanceA = simpleAMMHandler
            .ghost_expectedBalanceTokenA(user);
        uint256 actualBalanceA = tokenA.balanceOf(user);
        assertEq(
            expectedBalanceA,
            actualBalanceA,
            "User balance inconsistency for Token A"
        );

        uint256 expectedBalanceB = simpleAMMHandler
            .ghost_expectedBalanceTokenB(user);
        uint256 actualBalanceB = tokenB.balanceOf(user);
        assertEq(
            expectedBalanceB,
            actualBalanceB,
            "User balance inconsistency for Token B"
        );
    }
}

Final result

I have run 8 of our invariant tests, and this is our result:

Running 8 tests for test/SimpleAMM.t.sol:SimpleAMMTest
[PASS] invariant_ConstantProduct() (runs: 256, calls: 3840, reverts: 3771)
[PASS] invariant_LiquidityAdditionRemovalConsistency() (runs: 256, calls: 3840, reverts: 3765)
[PASS] invariant_PositiveLiquidity() (runs: 256, calls: 3840, reverts: 3776)
[PASS] invariant_ReserveNonNegativity() (runs: 256, calls: 3840, reverts: 3763)
[PASS] invariant_Solvency() (runs: 256, calls: 3840, reverts: 3756)
[PASS] invariant_SwapAmountValidation() (runs: 256, calls: 3840, reverts: 3760)
[PASS] invariant_SwapRateFairness() (runs: 256, calls: 3840, reverts: 3761)
[PASS] invariant_UserBalanceConsistency() (runs: 256, calls: 3840, reverts: 3781)
Test result: ok. 8 passed; 0 failed; 0 skipped; finished in 563.67ms

They all pass, and this is a very successful work guys. Now you should be ready to take this idea and start applying it in your smart contracts or on those audits you’re participating in.

Was it useful? Help us to improve!

With your feedback, we can improve the letter. Click on a link to vote:

• 👍 That helped me. Thanks

• 😐 Meh - was ok.

• 👎 Not interesting to me.