Cheat codes not available to contracts deployed during invariant testing

[webthethird]

Component

Forge

Have you ensured that all of these are up to date?

• Foundry
• Foundryup

What version of Foundry are you on?

forge 0.2.0 (1978a03 2023-12-21T09:45:30.854629898Z)

What command(s) is the bug in?

forge test

Operating System

Linux

Describe the bug

When testing systems which involve multiple deployments of the same contract, it may be desirable to deploy additional copies during the tests. For invariant testing, if a contract is more complex it is often necessary to wrap it in a handler contract. So if we want to deploy another copy of the contract we must also deploy another handler, which may require access to cheat codes such as deal.

However, it appears that cheats are not available to contracts deployed after the setUp function. Some copies of the handler, those which are deployed during setUp, do have cheat code access, but others deployed later from the same code are prevented from behaving the same way. The expected behavior is that my handler contracts behave the same regardless of when they were deployed, unless deployment time has some bearing on the underlying contract's execution.

The only way around this issue that I could find was to pre-deploy additional copies of the handler in question in the setUp function, despite not needing them until their underlying contracts are deployed. However, doing so muddles up the setUp function, adds additional overhead to the process of adding a new deployment and limits the number of deployments I can have to however many I pre-deployed. The helper function I use for targeting a newly deployed handler goes from this:

    function addStakedToken(
        StakedToken newStakedToken
    ) external onlySafetyModuleHandler {
        IERC20 underlying = newStakedToken.getUnderlyingToken();
        StakedTokenHandler newStakedTokenHandler = new StakedTokenHandler(newStakedToken, stakers);

        // Add staked token and its handler to lists
        stakedTokens.push(newStakedToken);
        stakedTokenHandlers.push(newStakedTokenHandler);

        // Register new target and exclude base contracts
        targetContract(address(newStakedTokenHandler));
        excludeContract(address(newStakedToken));
        excludeContract(address(underlying));
    }

to this:

    function addStakedToken(
        StakedToken newStakedToken
    ) external onlySafetyModuleHandler {
        IERC20 underlying = newStakedToken.getUnderlyingToken();

        // Use pre-deployed StakedTokenHandler
        // since deploying a new one doesn't give it access to cheats
        StakedTokenHandler newStakedTokenHandler;
        // Two StakedTokenHandlers are already deployed in setUp()
        if (numStakedTokenHandlers == 2)
            newStakedTokenHandler = stakedTokenHandler3;
        else if (numStakedTokenHandlers == 3)
            newStakedTokenHandler = stakedTokenHandler4;
        else if (numStakedTokenHandlers == 4)
            newStakedTokenHandler = stakedTokenHandler5;
        else if (numStakedTokenHandlers == 5)
            newStakedTokenHandler = stakedTokenHandler6;
        else if (numStakedTokenHandlers == 6)
            newStakedTokenHandler = stakedTokenHandler7;
        else if (numStakedTokenHandlers == 7)
            newStakedTokenHandler = stakedTokenHandler8;
        else if (numStakedTokenHandlers == 8)
            newStakedTokenHandler = stakedTokenHandler9;
        else
           revert("too many staked token handlers");
        numStakedTokenHandlers++;
        newStakedTokenHandler.setStakedToken(newStakedToken);

        // Add staked token and its handler to lists
        stakedTokens.push(newStakedToken);
        stakedTokenHandlers.push(newStakedTokenHandler);

        // Register new target and exclude base contracts
        targetContract(address(newStakedTokenHandler));
        excludeContract(address(newStakedToken));
        excludeContract(address(underlying));
    }

For additional context, here is the output for the failing test, where 0x1aF7f588A501EA2B5bB3feeFA744892aA2CF00e6 is the address of the new handler contract deployed during SafetyModuleHandler.addStakingToken on the fourth line from the bottom:

[FAIL. Reason: no cheats available for 0x1aF7f588A501EA2B5bB3feeFA744892aA2CF00e6]
	[Sequence]
		sender=0xaa9ec8bb61FCa06CF98F1462A29D796C232ed90C addr=[test/invariant/handlers/StakedBPTHandler.sol:StakedBPTHandler]0x15cF58144EF33af1e14b5208015d11F9143E27b9 calldata=transfer(uint256,uint256,uint256) args=[74, 122605228535406998642888989747 [1.226e29], 1]
		sender=0x0000000000000000000000000000000000001442 addr=[test/invariant/handlers/StakedBPTHandler.sol:StakedBPTHandler]0x15cF58144EF33af1e14b5208015d11F9143E27b9 calldata=setStakedToken(address) args=[0x000000000000000000000000000000000000429D]
		sender=0x0000000000000000000000000000000000003677 addr=[test/invariant/handlers/StakedBPTHandler.sol:StakedBPTHandler]0x15cF58144EF33af1e14b5208015d11F9143E27b9 calldata=cooldown(uint256) args=[16598 [1.659e4]]
		sender=0x0000000000000000000000000000000000004FD0 addr=[test/invariant/handlers/StakedTokenHandler.sol:StakedTokenHandler]0xD6BbDE9174b1CdAa358d2Cf4D57D1a9F7178FBfF calldata=redeem(uint256,uint256) args=[12944235812609692675684774955741423858894346459344382611099114779489443672409 [1.294e76], 0]
		sender=0x00000000000000000000000000000000000040f3 addr=[test/invariant/handlers/SMRDHandler.sol:SMRDHandler]0x03A6a84cD762D9707A21605b548aaaB891562aAb calldata=claimRewardsFor(uint256) args=[115792089237316195423570985008687907853269984665640564039457584007913129639933 [1.157e77]]
		sender=0x756666696369656E742062616C616e636520666e addr=[test/invariant/handlers/StakedTokenHandler.sol:StakedTokenHandler]0xD6BbDE9174b1CdAa358d2Cf4D57D1a9F7178FBfF calldata=dealUnderlying(uint256,uint256) args=[963, 7907]
		sender=0x00000000000000000000000000000000691c25e9 addr=[test/invariant/handlers/StakedTokenHandler.sol:StakedTokenHandler]0xD6BbDE9174b1CdAa358d2Cf4D57D1a9F7178FBfF calldata=stake(uint256,uint256) args=[2, 2]
		sender=0x0000000000000000000000000000000000000711 addr=[test/invariant/handlers/SMRDHandler.sol:SMRDHandler]0x03A6a84cD762D9707A21605b548aaaB891562aAb calldata=claimRewardsFor(uint256) args=[12822145867761120285609415092628159210061252848434669764027 [1.282e58]]
		sender=0x0000000000000000000000000000000000001442 addr=[test/invariant/handlers/StakedTokenHandler.sol:StakedTokenHandler]0xD6BbDE9174b1CdAa358d2Cf4D57D1a9F7178FBfF calldata=setStakedToken(address) args=[0x0000000000000000000000000000000000000Fcf]
		sender=0x00000000000000000000000000000000000002D2 addr=[test/invariant/handlers/StakedBPTHandler.sol:StakedBPTHandler]0x15cF58144EF33af1e14b5208015d11F9143E27b9 calldata=setStakedToken(address) args=[0xCCbD087951B520255269D83d499b01084b68236C]
		sender=0x00000000000000000000000000000000000015ED addr=[test/invariant/handlers/SMRDHandler.sol:SMRDHandler]0x03A6a84cD762D9707A21605b548aaaB891562aAb calldata=registerPositions(uint256) args=[23746 [2.374e4]]
		sender=0x00000000000000000000000000000000000077C6 addr=[test/invariant/handlers/SafetyModuleHandler.sol:SafetyModuleHandler]0xA4AD4f68d0b91CFD19687c881e50f3A00242828c calldata=addStakingToken(string,string,uint256,uint256,uint256) args=["𐤙🪿)\u{11d31}/ዕ¥ܝQ%u𖼚*e�େ`.", "FGg!\u{1cf1b}`'", 856770223169212166644023694771404365082446174 [8.567e44], 115792089237316195423570985008687907853269984665640564039457584007913129639933 [1.157e77], 4074004294613378912123582946270 [4.074e30]]
		sender=0xcd7Afc7dd5E122764Df67800943Ff2Cd44B7e8B7 addr=[test/invariant/handlers/SMRDHandler.sol:SMRDHandler]0x03A6a84cD762D9707A21605b548aaaB891562aAb calldata=registerPositions(uint256) args=[963292291715611957590915446518321895136973753255582368229899179331943 [9.632e68]]
		sender=0x0000000000000000000000000000000000000a16 addr=[test/invariant/handlers/StakedTokenHandler.sol:StakedTokenHandler]0x1aF7f588A501EA2B5bB3feeFA744892aA2CF00e6 calldata=transfer(uint256,uint256,uint256) args=[20697 [2.069e4], 1000000 [1e6], 2034]
 invariantExchangeRates() (runs: 1, calls: 14, reverts: 1)

Note: looking at the traces from the failed test, I can't tell whether calls to vm.<cheatcode> work correctly, such as vm.startPrank. In the following trace snippet, which is the last call in the traces and the first call to the newly deployed handler, it appears that vm.startPrank, vm.stopPrank, vm.expectRevert and vm.assume executed as usual:

  [56150] StakedTokenHandler::transfer(20697 [2.069e4], 1000000 [1e6], 2034)
    ├─ [0] console::log("Bound Result", 1) [staticcall]
    │   └─ ← ()
    ├─ [0] VM::startPrank(0x00000000000000000000000000000000000001c8)
    │   └─ ← ()
    ├─ [0] console::log("Bound Result", 0) [staticcall]
    │   └─ ← ()
    ├─ [0] VM::assume(true) [staticcall]
    │   └─ ← ()
    ├─ [9847] StakedToken::paused() [staticcall]
    │   ├─ [2382] SafetyModule::paused() [staticcall]
    │   │   └─ ← false
    │   └─ ← false
    ├─ [2428] StakedToken::maxStakeAmount() [staticcall]
    │   └─ ← 844613378912123578831115 [8.446e23]
    ├─ [2607] StakedToken::balanceOf(0x000000000000000000000000000000000000007B) [staticcall]
    │   └─ ← 0
    ├─ [2607] StakedToken::balanceOf(0x00000000000000000000000000000000000001c8) [staticcall]
    │   └─ ← 0
    ├─ [0] VM::expectRevert(ERC20: transfer amount exceeds balance)
    │   └─ ← ()
    ├─ [6975] StakedToken::transfer(0x000000000000000000000000000000000000007B, 20697 [2.069e4])
    │   ├─ [382] SafetyModule::paused() [staticcall]
    │   │   └─ ← false
    │   └─ ← revert: ERC20: transfer amount exceeds balance
    ├─ [0] VM::stopPrank()
    │   └─ ← ()
    └─ ← ()

The only cheat code used in the handler contract which is not a call to vm is deal(address token, address to, uint256 give), which is only used in StakedTokenHandler.dealUnderlying. But I still get the same no cheats available for 0x1aF7f588A501EA2B5bB3feeFA744892aA2CF00e6 error after commenting that line out. And in another failed test with the same reason, and without deal commented out, I get the following trace for dealUnderlying at the end of the call sequence:

  [161282] StakedTokenHandler::dealUnderlying(1342, 10386 [1.038e4])
    ├─ [0] console::log("Bound Result", 0) [staticcall]
    │   └─ ← ()
    ├─ [311] StakedToken::getUnderlyingToken() [staticcall]
    │   └─ ← ERC20: [0x3381cD18e2Fb4dB236BF0525938AB6E43Db0440f]
    ├─ [0] console::log("Bound Result", 999999000000000000001343 [9.999e23]) [staticcall]
    │   └─ ← ()
    ├─ [2561] ERC20::balanceOf(0x000000000000000000000000000000000000007B) [staticcall]
    │   └─ ← 0
    ├─ [0] VM::record()
    │   └─ ← ()
    ├─ [561] ERC20::balanceOf(0x000000000000000000000000000000000000007B) [staticcall]
    │   └─ ← 0
    ├─ [0] VM::accesses(ERC20: [0x3381cD18e2Fb4dB236BF0525938AB6E43Db0440f])
    │   └─ ← [0xde31a920dbdd1f015b2a842f0275dc8dec6a82ff94d9b796a36f23c64a3c8332], []
    ├─ [0] VM::load(ERC20: [0x3381cD18e2Fb4dB236BF0525938AB6E43Db0440f], 0xde31a920dbdd1f015b2a842f0275dc8dec6a82ff94d9b796a36f23c64a3c8332) [staticcall]
    │   └─ ← 0x0000000000000000000000000000000000000000000000000000000000000000
    ├─ emit WARNING_UninitedSlot(who: ERC20: [0x3381cD18e2Fb4dB236BF0525938AB6E43Db0440f], slot: 100501195172664535830925975091664947533607165983939194257793843064591798272818 [1.005e77])
    ├─ emit SlotFound(who: ERC20: [0x3381cD18e2Fb4dB236BF0525938AB6E43Db0440f], fsig: 0x70a0823100000000000000000000000000000000000000000000000000000000, keysHash: 0xde31a920dbdd1f015b2a842f0275dc8dec6a82ff94d9b796a36f23c64a3c8332, slot: 100501195172664535830925975091664947533607165983939194257793843064591798272818 [1.005e77])
    ├─ [561] ERC20::balanceOf(0x000000000000000000000000000000000000007B) [staticcall]
    │   └─ ← 0
    ├─ [0] VM::load(ERC20: [0x3381cD18e2Fb4dB236BF0525938AB6E43Db0440f], 0xde31a920dbdd1f015b2a842f0275dc8dec6a82ff94d9b796a36f23c64a3c8332) [staticcall]
    │   └─ ← 0x0000000000000000000000000000000000000000000000000000000000000000
    ├─ [0] VM::store(ERC20: [0x3381cD18e2Fb4dB236BF0525938AB6E43Db0440f], 0xde31a920dbdd1f015b2a842f0275dc8dec6a82ff94d9b796a36f23c64a3c8332, 0x00000000000000000000000000000000000000000000d3c20dee1639f99c053f)
    │   └─ ← ()
    └─ ← ()

With these traces, it is unclear which cheat code was unavailable.

[webthethird]

Are there any plans to fix this issue in the near future? Or should I abandon invariant testing with the deployment of new contracts for now?

[onbjerg]

@webthethird I assume you are forking in these tests. In this mode, we disallow cheatcode usage from anything other than the caller and the test contract as a safety precaution. You can allow other addresses to use cheats with vm.allowCheatcodes(address). Does this help?

[webthethird]

@onbjerg Yes I am forking, good catch. Your solution definitely seems promising! I'm working on something else at the moment, but I'll close this issue for now and re-open it if vm.allowCheatcodes doesn't work.
Thanks!