feat(cheatcodes): add vm.delegatePrank to execute a delegatecall in the context of the pranked address

[ncitron]

Component

Forge

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

• Foundry
• Foundryup

What version of Foundry are you on?

forge 0.1.0 (d08a59e 2022-03-01T00:28:36.621137610+00:00)

What command(s) is the bug in?

forge test

Operating System

Linux

Describe the bug

When attempting to perform a delegate call from a test while a prank is enabled, it does not execute in the context of the pranked address. This likely is not the expected behavior for users.

[mds1]

This current behavior should probably be the expected behavior, right? If you call vm.prank(user) followed by address(x).delegatecall(...), and the delegatecall executes in the context of user, this enables delegatecalling from an EOA which is probably not something we should support since it's not valid EVM behavior.

You could workaround this by having the cheatcode throw an error if the next call after prank is a delegatecall and the specified user has no code. But could you expand on the use case for delegatecalling from the test contract? That might help figure out the best way to resolve this issue

[ncitron]

So actually the usecase I had in mind was for mocking the behavior of a smart contract on a mainnet fork.

I was recently working on the payloads for a governance proposal on Aave, and was coming up with a testing plan. The obvious choice was to spin up a mainnet fork, prank an account that has a lot of AAVE, and go through the normal process of making a proposal, voting on it, and warping to a time where I can execute it. However, I wanted a bit of a simpler process, and since my proposal was simply going to have the Aave governance smart contract delegatecall into a custom "payload" contract which then handles all the proposal logic (this is a common pattern in on-chain governance). I was hoping to prank into the aave governance contract address, and delegatecall into my payload contract myself so I didn't have to walk through the entire proposal process.

This obviously isn't too big of a deal. The first method of walking through the governance process normally still works (and is probably safer since its closer to what will really happen), and I could also just use the etch cheatcode to overwrite the aave gov contracts code and replace it with a contract that just delegatecalls into my payload.

[mds1]

Ah, I think I misunderstood what you were trying to do. Let me know if this is the correct explanation of the issue:

• MyTestContract is DSTest has a test which calls vm.prank(aaveGovernor), where aaveGovernor is a proxy that delegates to an arbitrary aavePayload contract specified later
• After prank, the test then calls aaveGovernor.execute(proposalId, aavePayload) which delegatecalls to aavePayload
• msg.sender in that delegatecall is not aaveGovernor but is actually MyTestContract

I don't know if that function sig / flow is exactly correct, but if that is representative of the problem then I agree that seems like a bug (contrary to my previous comment where I thought you were calling delegatecall directly from the test contract)

[ncitron]

Hmm this is not quite the flow I was using. It looks more like this:

• MyTestContract is DSTest has a test which calls vm.prank(aaveGovernor), where aaveGovernor is a proxy that delegates to an arbitrary aavePayload contract specified later
• After prank, the test then calls address(aavePayload).delegatecall(abi.encodeWithSig("execute()"))
• address(this) inside of aavePayload is still the test contract, instead of aaveGovernor

The expectation I had was that calling prank(aaveGovernor) then delegatecalling to aavePayload immediately after in my test would allow me to mock a delegatecall from aaveGovernor to aavePayload

[mds1]

Got it, makes sense. Agreed the current behavior is confusing, so I think there's two ways to improve it, let me know if you have other ideas:

1. Leave behavior as is, but throw an error if the next call after vm.prank is a delegatecall instead of a regular call. In this case a workaround would be your suggestion to use vm.etch to overwrite the aave governor with a contract that just delegatecalls
2. Update prank so that you can use it for delegatecalling from a test contract, but throw an error if the address passed to vm.prank(addr) before a delegatecall has no code (to ensure you can't delegatecall from an EOA). Is there any related change to how pranking tx.origin is impacted? I don't think anything around tx.origin needs to change, but just making sure

IMO your use case is valid and potentially common so I'd support the second option

[adhusson]

+1, would also be useful to me, the general case seems to be:

• gov-like contract wants to delegatecall a proposal/spell/whatever you call it
• to test this, prank a delegatecall from gov-contract to proposal contract
• run tests

[adhusson]

There's now a delegate-prank tool available here: https://github.com/adhusson/delegate-prank

edit: updated interface

// c will delegatecall dest.fn(args)
delegatePrank(c,address(dest),abi.encodeCall(fn,(args)));

[ckksec]

Actually the idea of pranking delegatecall() can be super helpful, especially when you easily (atleast easier) want to overwrite storage. From what I know, this had to be done by getting the storage slot. It got even more complicated when you wanted to overwrite structs etc. I made an example which overwrites the owner that is stored as a value mapped to a random number in a mapping of numbers => address (to show how it can make possible complex things easier). If there was an easier way to manipulate storage other than by getting the slots please let me know haha.
My example was inspired by adhusson.

There's now a delegate-prank tool available here: https://github.com/adhusson/delegate-prank

For any contract just do

Delegator d = addDelegation(address(myContract));
d.delegatecall(dest,abi.encodeCall(...));

(warning: it will replace your contract bytecode with a proxy + a delegator)

Your code pretty much does what's needed, however, it should also change back the contract code to the original. Here's an example on how I did that and how a prankDelegate() cheatcode might look like based on inspiration from your code (I don't know how exactly they're implemented though). It temporarily changes the code of the to be pranked contract to a proxy (as you did) but once the delegation is done, it changes it back to the source. Therefore the storage of the contract will be changed after the delegatecall but not the code. Also there's no need to create multiple proxies then.

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

import "forge-std/Test.sol";

/**
    This contract is deployed by the user.
    It has to have the exact same storage structure than the original.
    In this simple example we want to change the variable `owner[0]` in `AAveGovernor`.
    Previously, as far as I know, we'd have to get the storage slot and overwrite it.
    By delegatecalling it temporarily as a proxy to our copy, we can easily do this without the previous step.
 */
contract AaveGovernorCopy {
    mapping(uint => address) public owners;

    // This function does not exist on AaveGovernor
    function changeOwner(address _newOwner) external {
        owners[19193984145157575157191571358] = _newOwner;
    }
}

/**
    This contract is the "main" contract where we want to overwrite the ownership.
    Our goal is to easily (or easier) change the owner even though `AaveGovernor` doesn't have a function to do that.
    To make things more complicated, the owner isn't stored in a simple address state variable, but a nice mapping.
    And to make it more complicated, it's assigned to an annoying number.
 */
contract AaveGovernor {
    mapping(uint => address) public owners;
    
    constructor() {
        owners[19193984145157575157191571358] = address(1);
    }
}

/**
    This is the temporary proxy contract that will replace the main contract (`AAveGovernor`) temporarily.
 */
contract PrankDelegator {
    function prankDelegate(address _to, bytes memory _data) external returns(bool _success, bytes memory _returnData) {
        return _to.delegatecall(_data);
    }
}

// Our test contract
contract ContractTest is Test {
    AaveGovernor aaveGovernor = new AaveGovernor();
    AaveGovernorCopy aaveGovernorCopy = new AaveGovernorCopy();
    PrankDelegator prankDelegator;

    // @dev Temporarily changes `_from` into a proxy which then delegatecalls to `_to` with `_data`.
    // @param _from The address to delegate from
    // @param _to The address to delegate to
    // @param _data The data to delegate with to `_to`
    function prankDelegate(address _from, address _to, bytes memory _data) internal returns(bool _success, bytes memory _returnData) {
        bytes memory sourceCode;
        sourceCode = _from.code; // Temporarily store the current code of `_from`.
        
        vm.etch(_from, address(prankDelegator).code); // "Turn" `_from` into a temporary proxy `PrankDelegator` by replacing its code.
        (_success, _returnData) = _from.call(abi.encodeWithSignature("prankDelegate(address,bytes)",_to,_data)); // Make `_from` delegate to `_to`.
        vm.etch(_from, sourceCode); // Restore the code of `_from` to its previous code before turning it into a proxy.
    }

    function setUp() public {

        // Create our three contracts.
        // Theoretically you could store the bytecode of `PrankDelegator` and replace the to be delegated contract with that so
        // that there's no need to deploy a new contract when doing tests or when a new cheatcode should be implemented.
        aaveGovernor = new AaveGovernor();
        aaveGovernorCopy = new AaveGovernorCopy();
        prankDelegator = new PrankDelegator();
        
        // No need to label the delegator as we won't call it but only delegatecall to it but from the main contract.
        vm.label(address(aaveGovernor), "Aave");
        vm.label(address(aaveGovernorCopy), "Fake Aave");
    }

    // Test changing the owner in `AaveGovernor` without using a prank but instead turning it into a temporary proxy.
    function testChangeOwner() public {
        // Check and log the current owner.
        require(aaveGovernor.owners(19193984145157575157191571358) == address(1), "Owner is not 1");
        console.logAddress(aaveGovernor.owners(19193984145157575157191571358));

        // Now do the `prankDelegate()` call.
        prankDelegate(address(aaveGovernor), address(aaveGovernorCopy), abi.encodeWithSignature("changeOwner(address)",0xDeaDbeefdEAdbeefdEadbEEFdeadbeEFdEaDbeeF));
        
        // Check and log the new owner.
        require(aaveGovernor.owners(19193984145157575157191571358) == 0xDeaDbeefdEAdbeefdEadbEEFdeadbeEFdEaDbeeF, "Owner is not 0xDeaDbeefdEAdbeefdEadbEEFdeadbeEFdEaDbeeF");
        console.logAddress(aaveGovernor.owners(19193984145157575157191571358));
    }
}

Simple result:

More complex result with proof that AAveGovernor did the delegatecall:

[adhusson]

@ckksec Nice, I went a step further and now the bytecode is swapped back right before the contract does the delegatecall.

[arr00]

This feature is 100% needed. The delegate caller being the test is unexpected and useless.

[EdwardJES]

Hey @zerosnacks if this is available, I'd like to take it, please.

Could you please answer these two questions for me and provide any other starters or tips?

1. What would be some general starters for the signature of vm.delegate rank? vm.delegatePrank(addr, codeAtAddr, ....)

2. Would this likely involve something similar to the prank API?

   foundry/crates/cheatcodes/src/evm/prank.rs

   Line 6 in c8db1e4

   pub struct Prank {

Thank you.

[zerosnacks]

Hi @EdwardJES,

I think the proposed API would be to update vm.prank but there is some complexity in how to handle the additional parameter (as txOrigin already exists as second address parameter). This is a very commonly used cheatcode so we have to avoid making breaking changes. This could mean adding an additional parameter after txOrigin rather than a new vm.delegatePrank cheatcode.

This is the related prank implementation:

foundry/crates/cheatcodes/src/evm/prank.rs

Lines 47 to 73 in 96105b4

	impl Cheatcode for prank_0Call {
	fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
	let Self { msgSender } = self;
	prank(ccx, msgSender, None, true)
	}
	}
	impl Cheatcode for startPrank_0Call {
	fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
	let Self { msgSender } = self;
	prank(ccx, msgSender, None, false)
	}
	}
	impl Cheatcode for prank_1Call {
	fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
	let Self { msgSender, txOrigin } = self;
	prank(ccx, msgSender, Some(txOrigin), true)
	}
	}
	impl Cheatcode for startPrank_1Call {
	fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
	let Self { msgSender, txOrigin } = self;
	prank(ccx, msgSender, Some(txOrigin), false)
	}
	}

It would involve modifying the prank API.

From @mds1

Update prank so that you can use it for delegatecalling from a test contract, but throw an error if the address passed to vm.prank(addr) before a delegatecall has no code (to ensure you can't delegatecall from an EOA). Is there any related change to how pranking tx.origin is impacted? I don't think anything around tx.origin needs to change, but just making sure

It may require some experimentation so having an early draft could be beneficial for feedback.

[EdwardJES]

Hey @zerosnacks, I've compiled an early draft here at #8863. I thought I'd at least validate this logic before spinning up a PR.

(@mds1 you may also be interested )

Ty 🙏

[leovct]

+1

My use case is not that meaningful but I was trying to solve Ethernaut CTF level 6 which requires to exploit delegatecall to override the storage value of the main contract. However, the exploit does not work because the pranked address is not used when making the delegate call.