From d60ab2fb35afd654749b14c9721bc260af669940 Mon Sep 17 00:00:00 2001 From: Rajath Alex Date: Thu, 18 Apr 2024 14:14:53 -0400 Subject: [PATCH] feat: account multicall module (#100) **Motivation:** We need a Rewards Claimer that allows multiple rewards to be claimed with the Llama account as the caller. This can be achieved through a generic `Account Multicall` module. **Modifications:** * `LlamaAccountMulticallFactory` - To deploy the account multicall modules. * `LlamaBaseAccountExtension` - Base Account Extension contract with `onlyDelegateCall` modifier * `LlamaAccountMulticallGuard` - Guard on top of Llama Account's execute function. * `LlamaAccountMulticallExtension` - A multicall extension for the Llama Account with authorized targets and selectors * `LlamaAccountMulticallStorage` - Account Multicall Storage contract (To prevent storage collision with Llama Account while being delegate-called) * Account Multicall deploy scripts and input config. * Tests **Result:** `Account Multicall` module. **Gas Report:** `forge test --match-test=test_Multicall --gas-report` Screen Shot 2024-04-11 at 9 44 04 PM --- justfile | 21 ++- .../DeployLlamaAccountMulticallFactory.s.sol | 21 +++ .../DeployLlamaAccountMulticallModule.s.sol | 74 +++++++++ .../input/1/mockAccountMulticallConfig.json | 13 ++ .../11155111/accountMulticallConfig.json | 13 ++ .../LlamaAccountMulticallExtension.sol | 59 +++++++ .../LlamaAccountMulticallFactory.sol | 66 ++++++++ .../LlamaAccountMulticallGuard.sol | 44 +++++ .../LlamaAccountMulticallStorage.sol | 93 +++++++++++ src/common/LlamaBaseAccountExtension.sol | 23 +++ src/interfaces/ILlamaAccount.sol | 19 +++ src/interfaces/ILlamaCore.sol | 4 + test/LlamaPeripheryTestSetup.sol | 4 + .../LlamaAccountMulticallExtension.t.sol | 142 ++++++++++++++++ .../LlamaAccountMulticallFactory.t.sol | 130 +++++++++++++++ .../LlamaAccountMulticallGuard.t.sol | 96 +++++++++++ .../LlamaAccountMulticallStorage.t.sol | 67 ++++++++ .../LlamaAccountMulticallTestSetup.sol | 152 ++++++++++++++++++ test/mock/MockRewardsContract.sol | 38 +++++ 19 files changed, 1077 insertions(+), 2 deletions(-) create mode 100644 script/DeployLlamaAccountMulticallFactory.s.sol create mode 100644 script/DeployLlamaAccountMulticallModule.s.sol create mode 100644 script/input/1/mockAccountMulticallConfig.json create mode 100644 script/input/11155111/accountMulticallConfig.json create mode 100644 src/account-multicall/LlamaAccountMulticallExtension.sol create mode 100644 src/account-multicall/LlamaAccountMulticallFactory.sol create mode 100644 src/account-multicall/LlamaAccountMulticallGuard.sol create mode 100644 src/account-multicall/LlamaAccountMulticallStorage.sol create mode 100644 src/common/LlamaBaseAccountExtension.sol create mode 100644 test/account-multicall/LlamaAccountMulticallExtension.t.sol create mode 100644 test/account-multicall/LlamaAccountMulticallFactory.t.sol create mode 100644 test/account-multicall/LlamaAccountMulticallGuard.t.sol create mode 100644 test/account-multicall/LlamaAccountMulticallStorage.t.sol create mode 100644 test/account-multicall/LlamaAccountMulticallTestSetup.sol create mode 100644 test/mock/MockRewardsContract.sol diff --git a/justfile b/justfile index 0fa9749..09ab40e 100644 --- a/justfile +++ b/justfile @@ -18,14 +18,31 @@ run-script script_name flags='' sig='' args='': -vvvv {{flags}} mv _test test -run-deploy-voting-module-script flags: (run-script 'DeployLlamaTokenVotingModule' flags '--sig "run(address,string)"' '$SCRIPT_DEPLOYER_ADDRESS "tokenVotingModuleConfig.json"') +# Token voting module dry-run-deploy: (run-script 'DeployLlamaTokenVotingFactory') deploy: (run-script 'DeployLlamaTokenVotingFactory' '--broadcast --verify --slow --build-info --build-info-path build_info') +verify: (run-script 'DeployLlamaTokenVotingFactory' '--verify --resume') + +run-deploy-voting-module-script flags: (run-script 'DeployLlamaTokenVotingModule' flags '--sig "run(address,string)"' '$SCRIPT_DEPLOYER_ADDRESS "tokenVotingModuleConfig.json"') + dry-run-deploy-voting-module: (run-deploy-voting-module-script '') deploy-voting-module: (run-deploy-voting-module-script '--broadcast --verify') -verify: (run-script 'DeployLlamaTokenVotingFactory' '--verify --resume') +# Account multicall module + +dry-run-deploy-account-multicall-factory: (run-script 'DeployLlamaAccountMulticallFactory') + +deploy-account-multicall-factory: (run-script 'DeployLlamaAccountMulticallFactory' '--broadcast --verify --slow --build-info --build-info-path build_info') + +verify-account-multicall-factory: (run-script 'DeployLlamaAccountMulticallFactory' '--verify --resume') + +run-deploy-account-multicall-module-script flags: (run-script 'DeployLlamaAccountMulticallModule' flags '--sig "run(address,string)"' '$SCRIPT_DEPLOYER_ADDRESS "accountMulticallConfig.json"') + +dry-run-deploy-account-multicall-module: (run-deploy-account-multicall-module-script '') + +deploy-account-multicall-module: (run-deploy-account-multicall-module-script '--broadcast --verify --slow') + diff --git a/script/DeployLlamaAccountMulticallFactory.s.sol b/script/DeployLlamaAccountMulticallFactory.s.sol new file mode 100644 index 0000000..f216f65 --- /dev/null +++ b/script/DeployLlamaAccountMulticallFactory.s.sol @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.23; + +import {Script, stdJson} from "forge-std/Script.sol"; + +import {DeployUtils} from "script/DeployUtils.sol"; + +import {LlamaAccountMulticallFactory} from "src/account-multicall/LlamaAccountMulticallFactory.sol"; + +contract DeployLlamaAccountMulticallFactory is Script { + // Factory contracts. + LlamaAccountMulticallFactory accountMulticallFactory; + + function run() public { + DeployUtils.print(string.concat("Deploying Llama account multicall factory to chain:", vm.toString(block.chainid))); + + vm.broadcast(); + accountMulticallFactory = new LlamaAccountMulticallFactory(); + DeployUtils.print(string.concat(" LlamaAccountMulticallFactory: ", vm.toString(address(accountMulticallFactory)))); + } +} diff --git a/script/DeployLlamaAccountMulticallModule.s.sol b/script/DeployLlamaAccountMulticallModule.s.sol new file mode 100644 index 0000000..7998652 --- /dev/null +++ b/script/DeployLlamaAccountMulticallModule.s.sol @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.23; + +import {Script, stdJson} from "forge-std/Script.sol"; + +import {DeployUtils} from "script/DeployUtils.sol"; + +import {LlamaAccountMulticallExtension} from "src/account-multicall/LlamaAccountMulticallExtension.sol"; +import {LlamaAccountMulticallFactory} from "src/account-multicall/LlamaAccountMulticallFactory.sol"; +import {LlamaAccountMulticallGuard} from "src/account-multicall/LlamaAccountMulticallGuard.sol"; +import {LlamaAccountMulticallStorage} from "src/account-multicall/LlamaAccountMulticallStorage.sol"; + +contract DeployLlamaAccountMulticallModule is Script { + using stdJson for string; + + struct TargetSelectorAuthorizationInputs { + // Attributes need to be in alphabetical order so JSON decodes properly. + string comment; + bytes selector; + address target; + } + + // Account multicall contracts. + LlamaAccountMulticallStorage accountMulticallStorage; + LlamaAccountMulticallExtension accountMulticallExtension; + LlamaAccountMulticallGuard accountMulticallGuard; + + function run(address deployer, string memory configFile) public { + string memory jsonInput = DeployUtils.readScriptInput(configFile); + + LlamaAccountMulticallFactory factory = LlamaAccountMulticallFactory(jsonInput.readAddress(".factory")); + + DeployUtils.print(string.concat("Deploying Llama account multicall module to chain:", vm.toString(block.chainid))); + + address llamaExecutor = jsonInput.readAddress(".llamaExecutor"); + uint256 nonce = jsonInput.readUint(".nonce"); + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data = readTargetSelectorAuthorizations(jsonInput); + LlamaAccountMulticallFactory.LlamaAccountMulticallConfig memory config = + LlamaAccountMulticallFactory.LlamaAccountMulticallConfig(llamaExecutor, nonce, data); + + vm.broadcast(deployer); + (accountMulticallGuard, accountMulticallExtension, accountMulticallStorage) = factory.deploy(config); + + DeployUtils.print("Successfully deployed a new Llama account multicall module"); + DeployUtils.print(string.concat(" LlamaAccountMulticallGuard: ", vm.toString(address(accountMulticallGuard)))); + DeployUtils.print( + string.concat(" LlamaAccountMulticallExtension: ", vm.toString(address(accountMulticallExtension))) + ); + DeployUtils.print( + string.concat(" LlamaAccountMulticallStorage: ", vm.toString(address(accountMulticallStorage))) + ); + } + + function readTargetSelectorAuthorizations(string memory jsonInput) + internal + pure + returns (LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory) + { + bytes memory data = jsonInput.parseRaw(".initialTargetSelectorAuthorizations"); + + TargetSelectorAuthorizationInputs[] memory rawConfigs = abi.decode(data, (TargetSelectorAuthorizationInputs[])); + + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory configs = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](rawConfigs.length); + + for (uint256 i = 0; i < rawConfigs.length; i++) { + configs[i].target = rawConfigs[i].target; + configs[i].selector = bytes4(rawConfigs[i].selector); + configs[i].isAuthorized = true; + } + + return configs; + } +} diff --git a/script/input/1/mockAccountMulticallConfig.json b/script/input/1/mockAccountMulticallConfig.json new file mode 100644 index 0000000..e6f8648 --- /dev/null +++ b/script/input/1/mockAccountMulticallConfig.json @@ -0,0 +1,13 @@ +{ + "comment": "This is an account multicall deployment for Forge tests.", + "factory": "0x90193C961A926261B756D1E5bb255e67ff9498A1", + "llamaExecutor": "0xdAf00E9786cABB195a8a1Cf102730863aE94Dd75", + "nonce": 0, + "initialTargetSelectorAuthorizations": [ + { + "comment": "DeadBeef.withdraw()", + "selector": "0x3ccfd60b", + "target": "0x00000000000000000000000000000000deadbeef" + } + ] +} diff --git a/script/input/11155111/accountMulticallConfig.json b/script/input/11155111/accountMulticallConfig.json new file mode 100644 index 0000000..aa86c86 --- /dev/null +++ b/script/input/11155111/accountMulticallConfig.json @@ -0,0 +1,13 @@ +{ + "comment": "This is an example account multicall deployment on Sepolia.", + "factory": "0x0000000000000000000000000000000000000000", + "llamaExecutor": "0x0000000000000000000000000000000000000000", + "nonce": 0, + "initialTargetSelectorAuthorizations": [ + { + "comment": "Target::Selector", + "selector": "0x00000000", + "target": "0x0000000000000000000000000000000000000000" + } + ] +} diff --git a/src/account-multicall/LlamaAccountMulticallExtension.sol b/src/account-multicall/LlamaAccountMulticallExtension.sol new file mode 100644 index 0000000..92ea291 --- /dev/null +++ b/src/account-multicall/LlamaAccountMulticallExtension.sol @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.23; + +import {LlamaAccountMulticallStorage} from "src/account-multicall/LlamaAccountMulticallStorage.sol"; +import {LlamaBaseAccountExtension} from "src/common/LlamaBaseAccountExtension.sol"; +import {LlamaUtils} from "src/lib/LlamaUtils.sol"; + +/// @title Llama Account Multicall Extension +/// @author Llama (devsdosomething@llama.xyz) +/// @notice An account extension that can multicall on behalf of the Llama account. +/// @dev This contract should be delegatecalled from a Llama account. +contract LlamaAccountMulticallExtension is LlamaBaseAccountExtension { + /// @dev Struct to hold target data. + struct TargetData { + address target; // The target contract. + uint256 value; // The target call value. + bytes data; // The target call data. + } + + /// @dev The call did not succeed. + /// @param index Index of the target data being called. + /// @param revertData Data returned by the called function. + error CallReverted(uint256 index, bytes revertData); + + /// @dev Thrown if the target-selector is not authorized. + error UnauthorizedTargetSelector(address target, bytes4 selector); + + /// @notice The Llama account multicall storage contract. + LlamaAccountMulticallStorage public immutable ACCOUNT_MULTICALL_STORAGE; + + /// @dev Initializes the Llama account multicall extenstion. + constructor(LlamaAccountMulticallStorage accountMulticallStorage) { + ACCOUNT_MULTICALL_STORAGE = accountMulticallStorage; + } + + /// @notice Multicalls on behalf of the Llama account. + /// @param targetData The target data to multicall. + /// @return returnData The return data from the target calls. + function multicall(TargetData[] memory targetData) external onlyDelegateCall returns (bytes[] memory returnData) { + uint256 length = targetData.length; + returnData = new bytes[](length); + for (uint256 i = 0; i < length; i = LlamaUtils.uncheckedIncrement(i)) { + address target = targetData[i].target; + uint256 value = targetData[i].value; + bytes memory callData = targetData[i].data; + bytes4 selector = bytes4(callData); + + // Check if the target-selector is authorized. + if (!ACCOUNT_MULTICALL_STORAGE.authorizedTargetSelectors(target, selector)) { + revert UnauthorizedTargetSelector(target, selector); + } + + // Execute the call. + (bool success, bytes memory result) = target.call{value: value}(callData); + if (!success) revert CallReverted(i, result); + returnData[i] = result; + } + } +} diff --git a/src/account-multicall/LlamaAccountMulticallFactory.sol b/src/account-multicall/LlamaAccountMulticallFactory.sol new file mode 100644 index 0000000..5f2fc75 --- /dev/null +++ b/src/account-multicall/LlamaAccountMulticallFactory.sol @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.23; + +import {LlamaAccountMulticallExtension} from "src/account-multicall/LlamaAccountMulticallExtension.sol"; +import {LlamaAccountMulticallGuard} from "src/account-multicall/LlamaAccountMulticallGuard.sol"; +import {LlamaAccountMulticallStorage} from "src/account-multicall/LlamaAccountMulticallStorage.sol"; + +/// @title LlamaAccountMulticallFactory +/// @author Llama (devsdosomething@llama.xyz) +/// @notice This contract enables Llama instances to deploy an account multicall module. +contract LlamaAccountMulticallFactory { + /// @dev Configuration of new Llama account multicall module. + struct LlamaAccountMulticallConfig { + address llamaExecutor; // The address of the Llama executor. + uint256 nonce; // The nonce of the new account multicall module. + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] data; // The target-selectors to authorize. + } + + /// @dev Emitted when a new Llama account multicall module is created. + event LlamaAccountMulticallModuleCreated( + address indexed deployer, + address indexed llamaExecutor, + uint256 nonce, + address accountMulticallGuard, + address accountMulticallExtension, + address accountMulticallStorage, + uint256 chainId + ); + + /// @notice Deploys a new Llama account multicall module. + /// @param accountMulticallConfig The configuration of the new Llama account multicall module. + /// @return accountMulticallGuard The deployed account multicall guard. + /// @return accountMulticallExtension The deployed account multicall extension. + /// @return accountMulticallStorage The deployed account multicall storage. + function deploy(LlamaAccountMulticallConfig memory accountMulticallConfig) + external + returns ( + LlamaAccountMulticallGuard accountMulticallGuard, + LlamaAccountMulticallExtension accountMulticallExtension, + LlamaAccountMulticallStorage accountMulticallStorage + ) + { + bytes32 salt = + keccak256(abi.encodePacked(msg.sender, accountMulticallConfig.llamaExecutor, accountMulticallConfig.nonce)); + + // Deploy and initialize account multicall storage. + accountMulticallStorage = new LlamaAccountMulticallStorage{salt: salt}(accountMulticallConfig.llamaExecutor); + accountMulticallStorage.initializeAuthorizedTargetSelectors(accountMulticallConfig.data); + + // Deploy account multicall extension. + accountMulticallExtension = new LlamaAccountMulticallExtension{salt: salt}(accountMulticallStorage); + + // Deploy account multicall guard. + accountMulticallGuard = new LlamaAccountMulticallGuard{salt: salt}(accountMulticallExtension); + + emit LlamaAccountMulticallModuleCreated( + msg.sender, + accountMulticallConfig.llamaExecutor, + accountMulticallConfig.nonce, + address(accountMulticallGuard), + address(accountMulticallExtension), + address(accountMulticallStorage), + block.chainid + ); + } +} diff --git a/src/account-multicall/LlamaAccountMulticallGuard.sol b/src/account-multicall/LlamaAccountMulticallGuard.sol new file mode 100644 index 0000000..3487ae0 --- /dev/null +++ b/src/account-multicall/LlamaAccountMulticallGuard.sol @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.23; + +import {LlamaAccountMulticallExtension} from "src/account-multicall/LlamaAccountMulticallExtension.sol"; +import {ILlamaActionGuard} from "src/interfaces/ILlamaActionGuard.sol"; +import {ActionInfo} from "src/lib/Structs.sol"; + +/// @title Llama Account Multicall Guard +/// @author Llama (devsdosomething@llama.xyz) +/// @notice A guard that only allows the `LlamaAccountMulticallExtension.multicall` to be delegate-called +/// @dev This guard should be used to protect the `execute` function in the `LlamaAccount` contract +contract LlamaAccountMulticallGuard is ILlamaActionGuard { + /// @dev Thrown if the call is not authorized. + error UnauthorizedCall(address target, bytes4 selector, bool withDelegatecall); + + /// @notice The address of the Llama account multicall extension. + LlamaAccountMulticallExtension public immutable ACCOUNT_MULTICALL_EXTENSION; + + /// @dev Initializes the Llama account multicall guard. + constructor(LlamaAccountMulticallExtension accountMulticallExtension) { + ACCOUNT_MULTICALL_EXTENSION = accountMulticallExtension; + } + + /// @inheritdoc ILlamaActionGuard + function validateActionCreation(ActionInfo calldata actionInfo) external view { + // Decode the action calldata to get the LlamaAccount execute target, call type and call data. + (address target, bool withDelegatecall,, bytes memory data) = + abi.decode(actionInfo.data[4:], (address, bool, uint256, bytes)); + bytes4 selector = bytes4(data); + + // Check if the target is the Llama account multicall extension, selector is `multicall` and the call type is a + // delegatecall. + if ( + target != address(ACCOUNT_MULTICALL_EXTENSION) || selector != LlamaAccountMulticallExtension.multicall.selector + || !withDelegatecall + ) revert UnauthorizedCall(target, selector, withDelegatecall); + } + + /// @inheritdoc ILlamaActionGuard + function validatePreActionExecution(ActionInfo calldata actionInfo) external pure {} + + /// @inheritdoc ILlamaActionGuard + function validatePostActionExecution(ActionInfo calldata actionInfo) external pure {} +} diff --git a/src/account-multicall/LlamaAccountMulticallStorage.sol b/src/account-multicall/LlamaAccountMulticallStorage.sol new file mode 100644 index 0000000..140d4ae --- /dev/null +++ b/src/account-multicall/LlamaAccountMulticallStorage.sol @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.23; + +import {LlamaUtils} from "src/lib/LlamaUtils.sol"; + +/// @title Llama Account Multicall Storage +/// @author Llama (devsdosomething@llama.xyz) +/// @notice The storage contract for the `LlamaAccountMulticallExtension` contract. +/// @dev This is a separate storage contract to prevent storage collisions with the Llama account. +contract LlamaAccountMulticallStorage { + // ========================= + // ======== Structs ======== + // ========================= + + /// @dev Struct to hold authorized target-selectors. + struct TargetSelectorAuthorization { + address target; // The target contract. + bytes4 selector; // The selector of the function being called. + bool isAuthorized; // Is the target-selector authorized. + } + + // ======================== + // ======== Errors ======== + // ======================== + + /// @dev Thrown if `initializeAuthorizedTargetSelectors` is called again. + error AlreadyInitialized(); + + /// @dev Only callable by a Llama instance's executor. + error OnlyLlama(); + + // ======================== + // ======== Events ======== + // ======================== + + /// @notice Emitted when a target-selector is authorized. + event TargetSelectorAuthorized(address indexed target, bytes4 indexed selector, bool isAuthorized); + + // =================================== + // ======== Storage Variables ======== + // =================================== + + /// @notice The Llama instance's executor. + address public immutable LLAMA_EXECUTOR; + + /// @dev Whether the contract is initialized. + bool internal initialized; + + /// @notice Mapping of all authorized target-selectors. + mapping(address target => mapping(bytes4 selector => bool isAuthorized)) public authorizedTargetSelectors; + + // ====================================================== + // ======== Contract Creation and Initialization ======== + // ====================================================== + + /// @dev Sets the Llama executor. + constructor(address llamaExecutor) { + LLAMA_EXECUTOR = llamaExecutor; + } + + /// @notice Initializes the authorized target-selectors. + /// @dev This function can only be called once. It should be called as part of the contract deployment. + /// @param data The target-selectors to authorize. + function initializeAuthorizedTargetSelectors(TargetSelectorAuthorization[] memory data) external { + if (initialized) revert AlreadyInitialized(); + initialized = true; + _setAuthorizedTargetSelectors(data); + } + + // ================================ + // ======== External Logic ======== + // ================================ + + /// @notice Sets the authorized target-selectors. + /// @param data The target-selectors to authorize. + function setAuthorizedTargetSelectors(TargetSelectorAuthorization[] memory data) external { + if (msg.sender != LLAMA_EXECUTOR) revert OnlyLlama(); + _setAuthorizedTargetSelectors(data); + } + + // ================================ + // ======== Internal Logic ======== + // ================================ + + /// @dev Sets the authorized target-selectors. + function _setAuthorizedTargetSelectors(TargetSelectorAuthorization[] memory data) internal { + uint256 length = data.length; + for (uint256 i = 0; i < length; i = LlamaUtils.uncheckedIncrement(i)) { + authorizedTargetSelectors[data[i].target][data[i].selector] = data[i].isAuthorized; + emit TargetSelectorAuthorized(data[i].target, data[i].selector, data[i].isAuthorized); + } + } +} diff --git a/src/common/LlamaBaseAccountExtension.sol b/src/common/LlamaBaseAccountExtension.sol new file mode 100644 index 0000000..e4126fe --- /dev/null +++ b/src/common/LlamaBaseAccountExtension.sol @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.23; + +/// @dev This account extension is a template for creating new account extensions, and should not be used directly. +abstract contract LlamaBaseAccountExtension { + /// @dev Thrown if you try to CALL a function that has the `onlyDelegatecall` modifier. + error OnlyDelegateCall(); + + /// @dev Add this to your account extension's methods to ensure the account extension can only be used via + /// delegatecall, and not a regular call. + modifier onlyDelegateCall() { + if (address(this) == SELF) revert OnlyDelegateCall(); + _; + } + + /// @dev Address of the account extension contract. We save it off because during a delegatecall `address(this)` + /// refers to the caller's address, not this account extension's address. + address internal immutable SELF; + + constructor() { + SELF = address(this); + } +} diff --git a/src/interfaces/ILlamaAccount.sol b/src/interfaces/ILlamaAccount.sol index c9d059a..57d47ca 100644 --- a/src/interfaces/ILlamaAccount.sol +++ b/src/interfaces/ILlamaAccount.sol @@ -5,6 +5,10 @@ pragma solidity ^0.8.23; /// @author Llama (devsdosomething@llama.xyz) /// @notice This is the interface for Llama accounts which can be used to hold assets for a Llama instance. interface ILlamaAccount { + /// @dev External call failed. + /// @param result Data returned by the called function. + error FailedExecution(bytes result); + // -------- For Inspection -------- /// @notice Returns the address of the Llama instance's executor. @@ -20,4 +24,19 @@ interface ILlamaAccount { /// @return This return statement must be hardcoded to `true` to ensure that initializing an EOA /// (like the zero address) will revert. function initialize(bytes memory config) external returns (bool); + + // -------- Generic Execution -------- + + /// @notice Execute arbitrary calls from the Llama Account. + /// @dev Be careful and intentional while assigning permissions to a policyholder that can create an action to call + /// this function, especially while using the delegatecall functionality as it can lead to arbitrary code execution in + /// the context of this Llama account. + /// @param target The address of the contract to call. + /// @param withDelegatecall Whether to use delegatecall or call. + /// @param value The amount of ETH to send with the call, taken from the Llama Account. + /// @param callData The calldata to pass to the contract. + /// @return The result of the call. + function execute(address target, bool withDelegatecall, uint256 value, bytes calldata callData) + external + returns (bytes memory); } diff --git a/src/interfaces/ILlamaCore.sol b/src/interfaces/ILlamaCore.sol index c83dda5..4e22981 100644 --- a/src/interfaces/ILlamaCore.sol +++ b/src/interfaces/ILlamaCore.sol @@ -29,6 +29,10 @@ interface ILlamaCore { /// @param current The current state of the action. error InvalidActionState(ActionState current); + /// @dev Action execution failed. + /// @param reason Data returned by the function called by the action. + error FailedActionExecution(bytes reason); + function actionGuard(address target, bytes4 selector) external view returns (address guard); function actionsCount() external view returns (uint256); diff --git a/test/LlamaPeripheryTestSetup.sol b/test/LlamaPeripheryTestSetup.sol index 54f0392..a86f73d 100644 --- a/test/LlamaPeripheryTestSetup.sol +++ b/test/LlamaPeripheryTestSetup.sol @@ -36,6 +36,10 @@ contract LlamaPeripheryTestSetup is Test { address coreTeam4 = 0x6b45E38c87bfCa15ee90AAe2AFe3CFC58cE08F75; address coreTeam5 = 0xbdfcE43E5D2C7AA8599290d940c9932B8dBC94Ca; + // This is the address that we're using to deploy modules. It could be + // replaced with any address that we hold the private key for. + address LLAMA_MODULE_DEPLOYER = 0x3d9fEa8AeD0249990133132Bb4BC8d07C6a8259a; + // Mock protocol for action targets. MockProtocol public mockProtocol; diff --git a/test/account-multicall/LlamaAccountMulticallExtension.t.sol b/test/account-multicall/LlamaAccountMulticallExtension.t.sol new file mode 100644 index 0000000..75b50c4 --- /dev/null +++ b/test/account-multicall/LlamaAccountMulticallExtension.t.sol @@ -0,0 +1,142 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.23; + +import {Test, console2} from "forge-std/Test.sol"; + +import {IERC20} from "@openzeppelin/token/ERC20/IERC20.sol"; + +import {MockRewardsContract} from "test/mock/MockRewardsContract.sol"; +import {LlamaAccountMulticallTestSetup} from "test/account-multicall/LlamaAccountMulticallTestSetup.sol"; + +import {ILlamaAccount} from "src/interfaces/ILlamaAccount.sol"; +import {ILlamaCore} from "src/interfaces/ILlamaCore.sol"; +import {ActionInfo} from "src/lib/Structs.sol"; +import {LlamaAccountMulticallExtension} from "src/account-multicall/LlamaAccountMulticallExtension.sol"; +import {LlamaAccountMulticallStorage} from "src/account-multicall/LlamaAccountMulticallStorage.sol"; +import {LlamaBaseAccountExtension} from "src/common/LlamaBaseAccountExtension.sol"; + +contract LlamaAccountMulticallExtensionTest is LlamaAccountMulticallTestSetup { + event ethWithdrawn(address indexed from, address indexed to, uint256 amount); + event erc20Withdrawn(IERC20 indexed token, address indexed from, address indexed to, uint256 amount); + + function setUp() public override { + LlamaAccountMulticallTestSetup.setUp(); + } +} + +contract Constructor is LlamaAccountMulticallExtensionTest { + function test_SetsAccountMulticallStorage() public { + assertEq(address(accountMulticallExtension.ACCOUNT_MULTICALL_STORAGE()), address(accountMulticallStorage)); + } +} + +contract Multicall is LlamaAccountMulticallExtensionTest { + function _createAndQueueActionClaimRewards() public returns (ActionInfo memory actionInfo) { + LlamaAccountMulticallExtension.TargetData[] memory targetData = _setupClaimRewardsData(); + bytes memory accountExtensionData = abi.encodeCall(LlamaAccountMulticallExtension.multicall, (targetData)); + bytes memory data = + abi.encodeCall(ILlamaAccount.execute, (address(accountMulticallExtension), true, 0, accountExtensionData)); + + // Create an action to claim rewards. + vm.prank(coreTeam1); + uint256 actionId = CORE.createAction(CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data, ""); + actionInfo = ActionInfo(actionId, coreTeam1, CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data); + + // Approve and queue the action. + vm.prank(coreTeam1); + CORE.castApproval(CORE_TEAM_ROLE, actionInfo, ""); + vm.prank(coreTeam2); + CORE.castApproval(CORE_TEAM_ROLE, actionInfo, ""); + vm.prank(coreTeam3); + CORE.castApproval(CORE_TEAM_ROLE, actionInfo, ""); + } + + function test_Multicall() public { + ActionInfo memory actionInfo = _createAndQueueActionClaimRewards(); + + assertEq(address(rewardsContract1).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract1)), 1 ether); + assertEq(address(rewardsContract2).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract2)), 1 ether); + assertEq(address(rewardsContract3).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract3)), 1 ether); + assertEq(address(rewardsContract4).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract4)), 1 ether); + assertEq(address(rewardsContract5).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract5)), 1 ether); + + uint256 initialLlamaAccountETHBalance = address(ACCOUNT).balance; + uint256 initialLlamaAccountUSDCBalance = USDC.balanceOf(address(ACCOUNT)); + + // Execute the action. + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract1), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract1), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract2), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract2), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract3), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract3), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract4), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract4), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract5), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract5), address(ACCOUNT), 1 ether); + + CORE.executeAction(actionInfo); + + assertEq(address(rewardsContract1).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract1)), 0); + assertEq(address(rewardsContract2).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract2)), 0); + assertEq(address(rewardsContract3).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract3)), 0); + assertEq(address(rewardsContract4).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract4)), 0); + assertEq(address(rewardsContract5).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract5)), 0); + + assertEq(address(ACCOUNT).balance, initialLlamaAccountETHBalance + 5 ether); + assertEq(USDC.balanceOf(address(ACCOUNT)), initialLlamaAccountUSDCBalance + 5 ether); + } + + function test_RevertIf_NotAuthorizedTargetSelector() public { + ActionInfo memory actionInfo = _createAndQueueActionClaimRewards(); + + // Unauthorize target selector. + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](1); + data[0] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract1), MockRewardsContract.withdrawETH.selector, false + ); + vm.prank(address(EXECUTOR)); + accountMulticallStorage.setAuthorizedTargetSelectors(data); + + bytes memory expectedErr = abi.encodeWithSelector( + ILlamaCore.FailedActionExecution.selector, + abi.encodeWithSelector( + ILlamaAccount.FailedExecution.selector, + abi.encodeWithSelector( + LlamaAccountMulticallExtension.UnauthorizedTargetSelector.selector, + address(rewardsContract1), + MockRewardsContract.withdrawETH.selector + ) + ) + ); + vm.expectRevert(expectedErr); + CORE.executeAction(actionInfo); + } + + function test_RevertIf_NotDelegateCalled() public { + LlamaAccountMulticallExtension.TargetData[] memory targetData = _setupClaimRewardsData(); + vm.expectRevert(LlamaBaseAccountExtension.OnlyDelegateCall.selector); + accountMulticallExtension.multicall(targetData); + } +} diff --git a/test/account-multicall/LlamaAccountMulticallFactory.t.sol b/test/account-multicall/LlamaAccountMulticallFactory.t.sol new file mode 100644 index 0000000..20cfb14 --- /dev/null +++ b/test/account-multicall/LlamaAccountMulticallFactory.t.sol @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.23; + +import {Test, console2} from "forge-std/Test.sol"; + +import {LlamaPeripheryTestSetup} from "test/LlamaPeripheryTestSetup.sol"; + +import {DeployLlamaAccountMulticallFactory} from "script/DeployLlamaAccountMulticallFactory.s.sol"; + +import {LlamaAccountMulticallExtension} from "src/account-multicall/LlamaAccountMulticallExtension.sol"; +import {LlamaAccountMulticallFactory} from "src/account-multicall/LlamaAccountMulticallFactory.sol"; +import {LlamaAccountMulticallGuard} from "src/account-multicall/LlamaAccountMulticallGuard.sol"; +import {LlamaAccountMulticallStorage} from "src/account-multicall/LlamaAccountMulticallStorage.sol"; + +contract LlamaAccountMulticallFactoryTest is LlamaPeripheryTestSetup, DeployLlamaAccountMulticallFactory { + event LlamaAccountMulticallModuleCreated( + address indexed deployer, + address indexed llamaExecutor, + uint256 nonce, + address accountMulticallGuard, + address accountMulticallExtension, + address accountMulticallStorage, + uint256 chainId + ); + + function setUp() public override { + LlamaPeripheryTestSetup.setUp(); + + // Deploy the factory + DeployLlamaAccountMulticallFactory.run(); + } + + // ========================= + // ======== Helpers ======== + // ========================= + + function _getCreate2Address(address deployer, uint256 salt, bytes memory bytecode) public pure returns (address) { + bytes32 hash = keccak256(abi.encodePacked(bytes1(0xff), deployer, salt, keccak256(bytecode))); + + // NOTE: cast last 20 bytes of hash to address + return address(uint160(uint256(hash))); + } + + function _preComputeMulticallAddresses(address moduleDeployer, uint256 nonce) + public + view + returns ( + address accountMulticallGuardAddress, + address accountMulticallExtensionAddress, + address accountMulticallStorageAddress + ) + { + bytes32 salt = keccak256(abi.encodePacked(moduleDeployer, address(EXECUTOR), nonce)); + + bytes memory multicallStorageBytecode = + abi.encodePacked(type(LlamaAccountMulticallStorage).creationCode, abi.encode(address(EXECUTOR))); + accountMulticallStorageAddress = + _getCreate2Address(address(accountMulticallFactory), uint256(salt), multicallStorageBytecode); + + bytes memory multicallExtensionBytecode = + abi.encodePacked(type(LlamaAccountMulticallExtension).creationCode, abi.encode(accountMulticallStorageAddress)); + accountMulticallExtensionAddress = + _getCreate2Address(address(accountMulticallFactory), uint256(salt), multicallExtensionBytecode); + + bytes memory multicallGuardBytecode = + abi.encodePacked(type(LlamaAccountMulticallGuard).creationCode, abi.encode(accountMulticallExtensionAddress)); + accountMulticallGuardAddress = + _getCreate2Address(address(accountMulticallFactory), uint256(salt), multicallGuardBytecode); + } +} + +contract DeployAccountMulticallModule is LlamaAccountMulticallFactoryTest { + function test_Deploy() public { + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](1); + data[0] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(0xdeadbeef), bytes4(keccak256("withdraw()")), true + ); + LlamaAccountMulticallFactory.LlamaAccountMulticallConfig memory config = + LlamaAccountMulticallFactory.LlamaAccountMulticallConfig(address(EXECUTOR), 0, data); + + ( + address accountMulticallGuardAddress, + address accountMulticallExtensionAddress, + address accountMulticallStorageAddress + ) = _preComputeMulticallAddresses(address(this), 0); + + vm.expectEmit(); + emit LlamaAccountMulticallModuleCreated( + address(this), + address(EXECUTOR), + 0, + accountMulticallGuardAddress, + accountMulticallExtensionAddress, + accountMulticallStorageAddress, + block.chainid + ); + ( + LlamaAccountMulticallGuard accountMulticallGuard, + LlamaAccountMulticallExtension accountMulticallExtension, + LlamaAccountMulticallStorage accountMulticallStorage + ) = accountMulticallFactory.deploy(config); + + assertEq(address(accountMulticallGuard.ACCOUNT_MULTICALL_EXTENSION()), address(accountMulticallExtension)); + assertEq(address(accountMulticallExtension.ACCOUNT_MULTICALL_STORAGE()), address(accountMulticallStorage)); + assertEq(address(accountMulticallStorage.LLAMA_EXECUTOR()), address(EXECUTOR)); + assertTrue(accountMulticallStorage.authorizedTargetSelectors(address(0xdeadbeef), bytes4(keccak256("withdraw()")))); + } + + function test_RevertIf_SameDeployerExecutorNonce() public { + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data1 = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](1); + data1[0] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(0xdeadbeef), bytes4(keccak256("withdraw()")), true + ); + LlamaAccountMulticallFactory.LlamaAccountMulticallConfig memory config = + LlamaAccountMulticallFactory.LlamaAccountMulticallConfig(address(EXECUTOR), 0, data1); + + // Initial deploy + accountMulticallFactory.deploy(config); + + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data2 = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](0); + config = LlamaAccountMulticallFactory.LlamaAccountMulticallConfig(address(EXECUTOR), 0, data2); + + // Revert on same deployer, executor, and nonce. Even if the TargetSelectorAuthorization data is different. + vm.expectRevert(); + accountMulticallFactory.deploy(config); + } +} diff --git a/test/account-multicall/LlamaAccountMulticallGuard.t.sol b/test/account-multicall/LlamaAccountMulticallGuard.t.sol new file mode 100644 index 0000000..004d64c --- /dev/null +++ b/test/account-multicall/LlamaAccountMulticallGuard.t.sol @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.23; + +import {Test, console2} from "forge-std/Test.sol"; + +import {LlamaAccountMulticallTestSetup} from "test/account-multicall/LlamaAccountMulticallTestSetup.sol"; + +import {ILlamaAccount} from "src/interfaces/ILlamaAccount.sol"; +import {LlamaAccountMulticallExtension} from "src/account-multicall/LlamaAccountMulticallExtension.sol"; +import {LlamaAccountMulticallGuard} from "src/account-multicall/LlamaAccountMulticallGuard.sol"; + +contract LlamaAccountMulticallGuardTest is LlamaAccountMulticallTestSetup { + function setUp() public override { + LlamaAccountMulticallTestSetup.setUp(); + } +} + +contract Constructor is LlamaAccountMulticallGuardTest { + function test_SetsAccountMulticallExtension() public { + assertEq(address(accountMulticallGuard.ACCOUNT_MULTICALL_EXTENSION()), address(accountMulticallExtension)); + } +} + +contract ValidateActionCreation is LlamaAccountMulticallGuardTest { + function test_PositiveFlow() public { + LlamaAccountMulticallExtension.TargetData[] memory targetData = _setupClaimRewardsData(); + + bytes memory accountExtensionData = abi.encodeCall(LlamaAccountMulticallExtension.multicall, (targetData)); + bytes memory data = + abi.encodeCall(ILlamaAccount.execute, (address(accountMulticallExtension), true, 0, accountExtensionData)); + + assertEq(CORE.actionsCount(), 12); + + vm.prank(coreTeam1); + CORE.createAction(CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data, ""); + + assertEq(CORE.actionsCount(), 13); + } + + function test_RevertIf_TargetIsNotAccountMulticallExtension() public { + LlamaAccountMulticallExtension.TargetData[] memory targetData = _setupClaimRewardsData(); + + address dummyTarget = address(0xdeadbeef); + + bytes memory accountExtensionData = abi.encodeCall(LlamaAccountMulticallExtension.multicall, (targetData)); + bytes memory data = abi.encodeCall(ILlamaAccount.execute, (dummyTarget, true, 0, accountExtensionData)); + + vm.prank(coreTeam1); + vm.expectRevert( + abi.encodeWithSelector( + LlamaAccountMulticallGuard.UnauthorizedCall.selector, + dummyTarget, + LlamaAccountMulticallExtension.multicall.selector, + true + ) + ); + CORE.createAction(CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data, ""); + } + + function test_RevertIf_SelectorIsNotMulticall() public { + LlamaAccountMulticallExtension.TargetData[] memory targetData = _setupClaimRewardsData(); + + bytes4 dummySelector = bytes4(0); + + bytes memory accountExtensionData = abi.encodeWithSelector(dummySelector, targetData); + bytes memory data = + abi.encodeCall(ILlamaAccount.execute, (address(accountMulticallExtension), true, 0, accountExtensionData)); + + vm.prank(coreTeam1); + vm.expectRevert( + abi.encodeWithSelector( + LlamaAccountMulticallGuard.UnauthorizedCall.selector, address(accountMulticallExtension), dummySelector, true + ) + ); + CORE.createAction(CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data, ""); + } + + function test_RevertIf_NotDelegateCall() public { + LlamaAccountMulticallExtension.TargetData[] memory targetData = _setupClaimRewardsData(); + + bytes memory accountExtensionData = abi.encodeCall(LlamaAccountMulticallExtension.multicall, (targetData)); + bytes memory data = + abi.encodeCall(ILlamaAccount.execute, (address(accountMulticallExtension), false, 0, accountExtensionData)); + + vm.prank(coreTeam1); + vm.expectRevert( + abi.encodeWithSelector( + LlamaAccountMulticallGuard.UnauthorizedCall.selector, + address(accountMulticallExtension), + LlamaAccountMulticallExtension.multicall.selector, + false + ) + ); + CORE.createAction(CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data, ""); + } +} diff --git a/test/account-multicall/LlamaAccountMulticallStorage.t.sol b/test/account-multicall/LlamaAccountMulticallStorage.t.sol new file mode 100644 index 0000000..73d28c5 --- /dev/null +++ b/test/account-multicall/LlamaAccountMulticallStorage.t.sol @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.23; + +import {Test, console2} from "forge-std/Test.sol"; + +import {LlamaAccountMulticallTestSetup} from "test/account-multicall/LlamaAccountMulticallTestSetup.sol"; + +import {LlamaAccountMulticallStorage} from "src/account-multicall/LlamaAccountMulticallStorage.sol"; + +contract LlamaAccountMulticallStorageTest is LlamaAccountMulticallTestSetup { + event TargetSelectorAuthorized(address indexed target, bytes4 indexed selector, bool isAuthorized); + + function setUp() public override { + LlamaAccountMulticallTestSetup.setUp(); + } +} + +contract Constructor is LlamaAccountMulticallStorageTest { + function test_SetsLlamaExecutor() public { + assertEq(address(accountMulticallStorage.LLAMA_EXECUTOR()), address(EXECUTOR)); + } +} + +contract InitializeAuthorizedTargetSelectors is LlamaAccountMulticallStorageTest { + function test_InitializeAuthorizedTargetSelectors() public { + assertTrue(accountMulticallStorage.authorizedTargetSelectors(address(0xdeadbeef), bytes4(keccak256("withdraw()")))); + } + + function test_RevertIf_AlreadyInitialized() public { + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](0); + + vm.prank(address(EXECUTOR)); + vm.expectRevert(LlamaAccountMulticallStorage.AlreadyInitialized.selector); + accountMulticallStorage.initializeAuthorizedTargetSelectors(data); + } +} + +contract SetAuthorizedTargetSelectors is LlamaAccountMulticallStorageTest { + function test_RevertIf_CallerIsNotLlama() public { + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](1); + data[0] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(0xdeadbeef), bytes4(keccak256("withdraw()")), false + ); + + vm.expectRevert(LlamaAccountMulticallStorage.OnlyLlama.selector); + accountMulticallStorage.setAuthorizedTargetSelectors(data); + } + + function test_SetAuthorizedTargetSelectors() public { + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](1); + data[0] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(0xdeadbeef), bytes4(keccak256("withdraw()")), false + ); + + assertTrue(accountMulticallStorage.authorizedTargetSelectors(address(0xdeadbeef), bytes4(keccak256("withdraw()")))); + + vm.prank(address(EXECUTOR)); + vm.expectEmit(); + emit TargetSelectorAuthorized(address(0xdeadbeef), bytes4(keccak256("withdraw()")), false); + accountMulticallStorage.setAuthorizedTargetSelectors(data); + + assertFalse(accountMulticallStorage.authorizedTargetSelectors(address(0xdeadbeef), bytes4(keccak256("withdraw()")))); + } +} diff --git a/test/account-multicall/LlamaAccountMulticallTestSetup.sol b/test/account-multicall/LlamaAccountMulticallTestSetup.sol new file mode 100644 index 0000000..dffba11 --- /dev/null +++ b/test/account-multicall/LlamaAccountMulticallTestSetup.sol @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.23; + +import {Test, console2} from "forge-std/Test.sol"; +import {Vm} from "forge-std/Vm.sol"; +import {StdUtils} from "forge-std/StdUtils.sol"; + +import {IERC20} from "@openzeppelin/token/ERC20/IERC20.sol"; + +import {LlamaPeripheryTestSetup} from "test/LlamaPeripheryTestSetup.sol"; +import {MockRewardsContract} from "test/mock/MockRewardsContract.sol"; + +import {DeployLlamaAccountMulticallFactory} from "script/DeployLlamaAccountMulticallFactory.s.sol"; +import {DeployLlamaAccountMulticallModule} from "script/DeployLlamaAccountMulticallModule.s.sol"; + +import {ILlamaAccount} from "src/interfaces/ILlamaAccount.sol"; +import {ILlamaPolicy} from "src/interfaces/ILlamaPolicy.sol"; +import {LlamaAccountMulticallExtension} from "src/account-multicall/LlamaAccountMulticallExtension.sol"; +import {LlamaAccountMulticallStorage} from "src/account-multicall/LlamaAccountMulticallStorage.sol"; + +contract LlamaAccountMulticallTestSetup is + LlamaPeripheryTestSetup, + DeployLlamaAccountMulticallFactory, + DeployLlamaAccountMulticallModule +{ + // Sample ERC20 token + IERC20 public constant USDC = IERC20(0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48); + + // Mock rewards contracts + MockRewardsContract rewardsContract1; + MockRewardsContract rewardsContract2; + MockRewardsContract rewardsContract3; + MockRewardsContract rewardsContract4; + MockRewardsContract rewardsContract5; + + function setUp() public virtual override { + LlamaPeripheryTestSetup.setUp(); + + // Deploy the factory + DeployLlamaAccountMulticallFactory.run(); + + // Deploy the Llama account multicall contracts. + DeployLlamaAccountMulticallModule.run(LLAMA_MODULE_DEPLOYER, "mockAccountMulticallConfig.json"); + + // Deploy the mock rewards contracts and set LlamaAccount as the reward claimer. + // Deal ETH and UNI to the rewards contracts. + rewardsContract1 = new MockRewardsContract(address(ACCOUNT)); + vm.deal(address(rewardsContract1), 1 ether); + deal(address(USDC), address(rewardsContract1), 1 ether); + + rewardsContract2 = new MockRewardsContract(address(ACCOUNT)); + vm.deal(address(rewardsContract2), 1 ether); + deal(address(USDC), address(rewardsContract2), 1 ether); + + rewardsContract3 = new MockRewardsContract(address(ACCOUNT)); + vm.deal(address(rewardsContract3), 1 ether); + deal(address(USDC), address(rewardsContract3), 1 ether); + + rewardsContract4 = new MockRewardsContract(address(ACCOUNT)); + vm.deal(address(rewardsContract4), 1 ether); + deal(address(USDC), address(rewardsContract4), 1 ether); + + rewardsContract5 = new MockRewardsContract(address(ACCOUNT)); + vm.deal(address(rewardsContract5), 1 ether); + deal(address(USDC), address(rewardsContract5), 1 ether); + + vm.startPrank(address(EXECUTOR)); + // Authorize the rewards contracts to claim rewards. + LlamaAccountMulticallStorage.TargetSelectorAuthorization[] memory data = + new LlamaAccountMulticallStorage.TargetSelectorAuthorization[](10); + data[0] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract1), MockRewardsContract.withdrawETH.selector, true + ); + data[1] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract1), MockRewardsContract.withdrawERC20.selector, true + ); + data[2] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract2), MockRewardsContract.withdrawETH.selector, true + ); + data[3] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract2), MockRewardsContract.withdrawERC20.selector, true + ); + data[4] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract3), MockRewardsContract.withdrawETH.selector, true + ); + data[5] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract3), MockRewardsContract.withdrawERC20.selector, true + ); + data[6] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract4), MockRewardsContract.withdrawETH.selector, true + ); + data[7] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract4), MockRewardsContract.withdrawERC20.selector, true + ); + data[8] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract5), MockRewardsContract.withdrawETH.selector, true + ); + data[9] = LlamaAccountMulticallStorage.TargetSelectorAuthorization( + address(rewardsContract5), MockRewardsContract.withdrawERC20.selector, true + ); + accountMulticallStorage.setAuthorizedTargetSelectors(data); + + // Set the Llama account multicall guard. + CORE.setGuard(address(ACCOUNT), ILlamaAccount.execute.selector, address(accountMulticallGuard)); + + // Assign LlamaAccount.execute permission to Core Team role. + POLICY.setRolePermission( + CORE_TEAM_ROLE, + ILlamaPolicy.PermissionData(address(ACCOUNT), ILlamaAccount.execute.selector, address(STRATEGY)), + true + ); + vm.stopPrank(); + } + + // ========================= + // ======== Helpers ======== + // ========================= + + function _setupClaimRewardsData() public view returns (LlamaAccountMulticallExtension.TargetData[] memory targetData) { + targetData = new LlamaAccountMulticallExtension.TargetData[](10); + targetData[0] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract1), 0, abi.encodeCall(MockRewardsContract.withdrawETH, ()) + ); + targetData[1] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract1), 0, abi.encodeCall(MockRewardsContract.withdrawERC20, (USDC)) + ); + targetData[2] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract2), 0, abi.encodeCall(MockRewardsContract.withdrawETH, ()) + ); + targetData[3] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract2), 0, abi.encodeCall(MockRewardsContract.withdrawERC20, (USDC)) + ); + targetData[4] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract3), 0, abi.encodeCall(MockRewardsContract.withdrawETH, ()) + ); + targetData[5] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract3), 0, abi.encodeCall(MockRewardsContract.withdrawERC20, (USDC)) + ); + targetData[6] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract4), 0, abi.encodeCall(MockRewardsContract.withdrawETH, ()) + ); + targetData[7] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract4), 0, abi.encodeCall(MockRewardsContract.withdrawERC20, (USDC)) + ); + targetData[8] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract5), 0, abi.encodeCall(MockRewardsContract.withdrawETH, ()) + ); + targetData[9] = LlamaAccountMulticallExtension.TargetData( + address(rewardsContract5), 0, abi.encodeCall(MockRewardsContract.withdrawERC20, (USDC)) + ); + } +} diff --git a/test/mock/MockRewardsContract.sol b/test/mock/MockRewardsContract.sol new file mode 100644 index 0000000..819365e --- /dev/null +++ b/test/mock/MockRewardsContract.sol @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.23; + +import {IERC20} from "@openzeppelin/token/ERC20/IERC20.sol"; +import {SafeERC20} from "@openzeppelin/token/ERC20/utils/SafeERC20.sol"; +import {Address} from "@openzeppelin/utils/Address.sol"; + +contract MockRewardsContract { + using SafeERC20 for IERC20; + + error OnlyRewardClaimer(); + + modifier onlyRewardClaimer() { + if (msg.sender != REWARD_CLAIMER) revert OnlyRewardClaimer(); + _; + } + + event ethWithdrawn(address indexed from, address indexed to, uint256 amount); + event erc20Withdrawn(IERC20 indexed token, address indexed from, address indexed to, uint256 amount); + + address public immutable REWARD_CLAIMER; + + constructor(address rewardClaimer) { + REWARD_CLAIMER = rewardClaimer; + } + + receive() external payable {} + + function withdrawETH() external payable onlyRewardClaimer { + emit ethWithdrawn(address(this), REWARD_CLAIMER, address(this).balance); + Address.sendValue(payable(REWARD_CLAIMER), address(this).balance); + } + + function withdrawERC20(IERC20 token) external onlyRewardClaimer { + emit erc20Withdrawn(token, address(this), REWARD_CLAIMER, token.balanceOf(address(this))); + token.safeTransfer(REWARD_CLAIMER, token.balanceOf(address(this))); + } +}