book icon indicating copy to clipboard operation
book copied to clipboard

feat(`prompt`): modify prompting guide to support codegen invariant handlers

Open pbshgthm opened this issue 4 years ago • 8 comments

Note: title updated

See: https://github.com/foundry-rs/book/issues/1573


User based testing pattern

A common pattern used in dapptools and foundry projects is a 'User Contract'. It is an abstraction over the contract interaction that looks something like this (from DSToken)

contract TokenUser {
    DSToken  token;

    constructor(DSToken token_) public {
        token = token_;
    }

    function doApprove(address spender, uint amount)
        public
        returns (bool)
    {
        return token.approve(spender, amount);
    }
    //...
}

And in the test contract, we can use it like this

user1 = new TokenUser(token);
user1.doApprove(user2);

This is an interesting pattern as it lets you wrap contracts and test them like how a user would interact with them. Its advantage shines really well when you need to test with multiple user addresses.

Blacksmiths to use the foundry

I took this pattern a step further and created a full-fledged contract generator that will create these 'User Contracts', along with a bunch of UX niceties. You can create a user with a particular address or private key and it can perform all operations an EOA can. All this can be automated by running the blacksmith.js script (sorry, not in rust yet). It automatically creates User contracts for all contracts in your foundry project directory.

Features

Wrap multiple target contracts

user1.dex.swap(100);
user1.factory.pause(true);
user1.token.transfer(user2.addr, 100);

Sign using private key

(uint8 v, bytes32 r, bytes32 s) = user1.sign("blacksmith");

Call arbiraty contracts

user1.call{value:10}(contract_address, "calldata");

Set user address's balance

user1.deal(100);

Zero code size at address

user1.addr.code.length // is zero

The blacksmith.js script creates the baseBlacksmith.sol that contains basic functions like call, sign and deal. It also creates TargetBS.sol for all Target contracts in the project directory. To Base User contract (Blacksmith) takes in an address and a private key as constructor params. If the private key is zero, the provided address is used as the user's address, else the address is calculated from the private key.

constructor( address _addr, uint256 _privateKey, address _target) {
    addr = _privateKey == 0 ? _addr : bsvm.addr(_privateKey);
    privateKey = _privateKey;
    target = payable(_target);
}

To create a User contract to interact with a Target contract, you import TargetBS contract. Along with the address and private key, it also takes in the target contract's address

constructor( address _addr, uint256 _privateKey, address _target) {
    addr = _privateKey == 0 ? _addr : bsvm.addr(_privateKey);
    privateKey = _privateKey;
    target = payable(_target);
}

To create a user object, you can create a struct and add the required interface. The below code is all you'll need to write to get started with testing. Rest is taken care of by blacksmith script.

struct User {
    address addr;  // to avoid external call, we save it in the struct
    Blacksmith base;  // contains call(), sign(), deal()
    FooTokenBS foo;  // interacts with FooToken contract
    BarTokenBS bar;  // interacts with BarToken contract
}

function createUser(address _addr, uint256 _privateKey) public returns (User memory) {
    Blacksmith base = new Blacksmith(_addr, _privateKey);
    FooTokenBS _foo = new FooTokenBS(_addr, _privateKey, address(foo));
    BarTokenBS _bar = new BarTokenBS(_addr, _privateKey, address(bar));
    base.deal(100);
    return User(base.addr(), base, _foo, _bar);
}

function setUp() public {
    foo = new FooToken();
    bar = new BarToken();
    alice = createUser(address(0), 111);  // addrss will be 0x052b91ad9732d1bce0ddae15a4545e5c65d02443
    bob = createUser(address(111), 0);  // address will be 0x000000000000000000000000000000000000006f
    eve = createUser(address(123), 0);  // address will be 0x000000000000000000000000000000000000006f
}

Now you can use it directly in your test functions

function testSomething() public {
    bob.foo.approve(alice.addr, 10);
    alice.foo.transferFrom(bob.addr, alice.addr, 10);
    alice.bar.approve(eve.addr, 100);
    eve.transferFrom(bob.addr, eve.addr, 50);
    eve.call{value:10}(alice.ddr, "");
   (uint8 v, bytes32 r, bytes32 s) = alice.sign("blacksmith");
}

Usage

To get started with blacksmith, download blacksmith.js to the foundry project’s root directory.

curl -O https://raw.githubusercontent.com/pbshgthm/blacksmith/main/blacksmith.js
node blacksmith.js create #in foundry project's root directory

This will run forge build and then create /src/test/blacksmith directory with user contracts in Target.bs.sol.

Would love to know what folks think about this pattern. Might be interesting to add it as a feature to forge as it would help avoid a ton of boilerplate code and prove UX improvements to testing contracts in an OOP like way.

Blacksmith Repo

pbshgthm avatar Feb 16 '22 15:02 pbshgthm

+1 on this being a great addition. I think that it might make sense that we start supporting the following structure:

src/ <-- all source contracts, no tests
test/ <-- test contracts
codegen/ <-- user contract bindings
lib/ <-- our deps

whereas now both blacksmith bindings and tests fall under src

this would prevent fatigue from having a src directory with "too many" things

supporting this kind of codegen in Rust is no prob

gakonst avatar Feb 16 '22 15:02 gakonst

Yes, this directory structure makes sense. Currently, any changes in the function signatures of the contract will result in compile error on forge build in the Blacksmith contracts, as they inherit from the target contract and function signatures won't' match. But this isn't a functional issue, as forge build compiles the rest of the contracts and the error can be fixed by running blacksmith.js. Still, it's something that can be improved. We can probably look at building only the src directory with forge build and other directories like test as required. So taking test out of src definitely seems like a good direction.

pbshgthm avatar Feb 16 '22 15:02 pbshgthm

Yeah we could do it in 2-steps with this separation: Build source -> run codegen -> build codegen.

gakonst avatar Feb 16 '22 15:02 gakonst

+1 on both this feature and the directory structure suggested by @gakonst

With native integration of this feature, compatibility with https://github.com/gakonst/foundry/issues/402 could be valuable. For example:

  • define a top level deploy() method that deploys your full protocol
  • each contract deployed in that method is automatically added to each blacksmith user object you create
  • that same deploy() method can be used for production deployments
  • that same deploy() method can be called during the setUp() method to run tests against that state

mds1 avatar Feb 16 '22 15:02 mds1

Revisiting this issue, I don't find myself ever using user contracts for testing, except when creating handler contracts during invariant testing. So maybe the best use of codegen contracts like this might be to simplify setup for invariant testing, to automatically generate handler contracts.

forge build could automatically write one handler contract for each contract in the src/ dir to a codegen directory, and a combined handler for everything. You can also imagine additional args to customize the handlers created in the config file. For example:

[handlers]
# Generates a contract called Handler A that can call all mutable methods on
# MyContract1 and MyContract2. `warp = true` means it will have a public
# method to skip forward in time, `roll = false` means it will *not* have a public
# method to skip forward in time. If both were true, it's a single method to modify both.
# Can also include things like tracking for call summaries and actors automaticallt
HandlerA = { contracts = ["MyContract1, MyContract2"], warp = true, roll = false }

Might need to think on it more, but I think this could very useful and more flexible than a Handler template in forge-std. cc @lucas-manuel

mds1 avatar Feb 25 '23 18:02 mds1

The codegen idea is really interesting. I think in practice handlers will want to have extra customized assertions in them, so probably hard to create a general enough solution.

gakonst avatar Mar 01 '23 01:03 gakonst

Ooh this could be cool. Have a codegen handler that imports from forge-std to use the useRandomActor modifier and has a wrapper that includes the fuzzed index and all the necessary params for wrapping the call. After that just calls function and leaves logic up to user.

lucas-manuel avatar Mar 12 '23 20:03 lucas-manuel

Updated the title to reflect the outcome of the discussion, curious @grandizzy if you see ways this could be relevant to simplify setting up of invariants as the boilerplate involved was initially quite a headache

zerosnacks avatar Jul 31 '24 15:07 zerosnacks