Open1

FoundryのInvariant Test

ぽけなぽけな

以下YouTubeのまとめ的なもの
https://www.youtube.com/watch?v=JtzBi67hgLI

Invariant Testとは

  • コントラクト内の関数をランダムに色々な順番で呼び出しまくっても異常な状態にならないかのテストができる
  • 以下の例では、func_1~5が様々な順番で呼ばれまくる
  • invariant_flag_is_always_falseではflagが常にfalseであることを期待しているが、func_5が呼ばれるとtrueになるためNGになる
contract InvariantIntro {
    bool public flag;

    function func_1() external {}
    function func_2() external {}
    function func_3() external {}
    function func_4() external {}

    function func_5() external {
        flag = true;
    }
}

contract IntroInvariantTest is Test {
    InvariantIntro private target;

    function setUp() public {
        target = new InvariantIntro();
    }

    function invariant_flag_is_always_false() public {
        assertEq(target.flag(), false);
    }
}

Invariant Test実行時の困る点

  • WETHのコントラクトにおいて、totalSupplyが常に0であることを試験しようとしている
  • 通常、WETHコントラクトにETHをdepositすればtotalSupplyはその分増えるし、withdrawすれば減るため、この試験は失敗するはず
  • しかし、この試験はOKになってしまう
    • 呼び出し元のコントラクトがETHを持っていないため、depositがrevertされてしまうため
    • withdrawもdepositをしていない(またはdeposit以上にwithdrawしようとする)ため、revertされてしまう
contract WETH_Open_Invariant_Tests is Test {
    WETH public weth;

    function setUp() public {
        weth = new WETH();
    }

    receive() external payable {}

    // NOTE: - calls = runs x depth, (runs, calls, reverts)
    function invariant_totalSupply_is_always_zero() public {
        assertEq(0, weth.totalSupply());
    }
}

上記問題の解決方法

  • Handlerと呼ばれるコントラクトを作成し、ETHの保持をしたり、depositした以上の金額をwithdrawできないよう制限する
    • setUp関数内のdealでHandlerコントラクトに100ETHを付与
    • depositやwithdraw関数を仲介してあげる関数を作成し、その内部のboundでrevertされない値を取るようにしている
  • テストコード→Handler→テスト対象のコントラクトの実行順になる
  • Handlerを使うようにするには、targetContract(address(handler));を実行する必要がある
  • Handler内の特定の関数しか呼びたくない場合には、Selectorを使う
contract Handler is CommonBase, StdCheats, StdUtils {
    WETH private weth;
    uint256 public wethBalance;
    uint256 public numCalls;

    constructor(WETH _weth) {
        weth = _weth;
    }

    receive() external payable {}

    function sendToFallback(uint256 amount) public {
        amount = bound(amount, 0, address(this).balance);
        wethBalance += amount;
        numCalls += 1;

        (bool ok,) = address(weth).call{value: amount}("");
        require(ok, "sendToFallback failed");
    }

    function deposit(uint256 amount) public {
        amount = bound(amount, 0, address(this).balance);
        wethBalance += amount;
        numCalls += 1;

        weth.deposit{value: amount}();
    }

    function withdraw(uint256 amount) public {
        amount = bound(amount, 0, weth.balanceOf(address(this)));
        wethBalance -= amount;
        numCalls += 1;

        weth.withdraw(amount);
    }

    function fail() external {
        revert("fail");
    }
}

contract WETH_Handler_Based_Invariant_Tests is Test {
    WETH public weth;
    Handler public handler;

    function setUp() public {
        weth = new WETH();
        handler = new Handler(weth);

        // Send 100 ETH to handler
        deal(address(handler), 100 * 1e18);
        // Set fuzzer to only call the handler
        targetContract(address(handler));

        bytes4[] memory selectors = new bytes4[](3);
        selectors[0] = Handler.deposit.selector;
        selectors[1] = Handler.withdraw.selector;
        selectors[2] = Handler.sendToFallback.selector;

        // Handler.fail() not called
        targetSelector(
            FuzzSelector({addr: address(handler), selectors: selectors})
        );
    }

    function invariant_eth_balance() public {
        assertGe(address(weth).balance, handler.wethBalance());
        console.log("handler num calls", handler.numCalls());
    }
}

複数のアドレスからのコールを確認したい場合

  • Managerの役割を持つコントラクトを作成し、複数のHandlerを使ってテストする
  • テスト関数内の引数はfuzzテスト時と同様にランダム値であるため、これを用いてランダムなHandlerを使ってテストができる
contract ActorManager is CommonBase, StdCheats, StdUtils {
    Handler[] public handlers;

    constructor(Handler[] memory _handlers) {
        handlers = _handlers;
    }

    function sendToFallback(uint256 handlerIndex, uint256 amount) public {
        uint256 index = bound(handlerIndex, 0, handlers.length - 1);
        handlers[index].sendToFallback(amount);
    }

    function deposit(uint256 handlerIndex, uint256 amount) public {
        uint256 index = bound(handlerIndex, 0, handlers.length - 1);
        handlers[index].deposit(amount);
    }

    function withdraw(uint256 handlerIndex, uint256 amount) public {
        uint256 index = bound(handlerIndex, 0, handlers.length - 1);
        handlers[index].withdraw(amount);
    }
}

contract WETH_Multi_Handler_Invariant_Tests is Test {
    WETH public weth;
    ActorManager public manager;
    Handler[] public handlers;

    function setUp() public {
        weth = new WETH();

        for (uint256 i = 0; i < 3; i++) {
            handlers.push(new Handler(weth));
            // Send 100 ETH to handler
            deal(address(handlers[i]), 100 * 1e18);
        }

        manager = new ActorManager(handlers);

        targetContract(address(manager));
    }

    function invariant_eth_balance() public {
        uint256 total = 0;
        for (uint256 i = 0; i < handlers.length; i++) {
            total += handlers[i].wethBalance();
            console.log("Handler num calls", i, handlers[i].numCalls());
        }
        console.log("ETH total", total / 1e18);
        assertGe(address(weth).balance, total);
    }
}

参考

https://github.com/t4sk/hello-foundry/tree/main/test/invariants
https://allthingsfuzzy.substack.com/p/creating-invariant-tests-for-an-amm