diff --git a/src/templates/MultisigBuilder.sol b/src/templates/MultisigBuilder.sol index beec8db..2cf7fdb 100644 --- a/src/templates/MultisigBuilder.sol +++ b/src/templates/MultisigBuilder.sol @@ -1,34 +1,48 @@ // SPDX-License-Identifier: BUSL-1.1 pragma solidity ^0.8.12; -import {ZeusScript} from "../utils/ZeusScript.sol"; +import "../utils/ZeusScript.sol"; /** * @title MultisigBuilder * @dev Abstract contract for building arbitrary multisig scripts. */ abstract contract MultisigBuilder is ZeusScript { + + bool private hasPranked; + + modifier prank(address caller) { + _startPrank(caller); + _; + _stopPrank(); + } + /** * @notice Constructs a SafeTx object for a Gnosis Safe to ingest. Emits via `ZeusMultisigExecute` */ function execute() public { - MultisigOptions memory opt = options(); - emit ZeusRequireMultisig(opt.addr, opt.callType); - vm.startPrank(opt.addr, opt.addr); - runAsMultisig(); - vm.stopPrank(); + _runAsMultisig(); + require(hasPranked, "MultisigBuilder.execute: did not use prank helpers"); } /** - * Indicate the multisig address + * @dev Implement the most high-level call performed by the target multisig for this script. + * Note: This function should be written as if the target multisig is performing the call directly. + * Note: This function MUST be written by using the `prank` modifier or `_startPrank`/`_stopPrank` + * helper methods. */ - function options() internal virtual returns (MultisigOptions memory); + function _runAsMultisig() internal virtual; - /** - * @notice To be implemented by inheriting contract. - * - * This function will be pranked from the perspective of the multisig you choose to run with. - * DO NOT USE vm.startPrank()/stopPrank() during your implementation. - */ - function runAsMultisig() internal virtual; + function _startPrank(address caller) internal { + require(!hasPranked, "MultisigBuilder._startPrank: called twice in txn"); + hasPranked = true; + + emit ZeusRequireMultisig(caller, Encode.Operation.Call); + vm.startPrank(caller); + } + + function _stopPrank() internal { + require(hasPranked, "MultisigBuilder._stopPrank: _startPrank not called"); + vm.stopPrank(); + } } diff --git a/src/utils/ZeusScript.sol b/src/utils/ZeusScript.sol index 60a65d3..a67efab 100644 --- a/src/utils/ZeusScript.sol +++ b/src/utils/ZeusScript.sol @@ -12,17 +12,12 @@ abstract contract ZeusScript is Script, Test { using StringUtils for string; using ZEnvHelpers for *; - enum Operation { - Call, - DelegateCall - } - struct MultisigOptions { address addr; // the address of the multisig - Operation callType; // call vs. delegateCall + Encode.Operation callType; // call vs. delegateCall } - event ZeusRequireMultisig(address addr, Operation callType); + event ZeusRequireMultisig(address addr, Encode.Operation callType); event ZeusEnvironmentUpdate(string key, EnvironmentVariableType internalType, bytes value); event ZeusDeploy(string name, address addr, bool singleton); event ZeusMultisigExecute(address to, uint256 value, bytes data, Encode.Operation op); diff --git a/test/ZeusScript.test.sol b/test/ZeusScript.test.sol index 993e83f..ccd962c 100644 --- a/test/ZeusScript.test.sol +++ b/test/ZeusScript.test.sol @@ -404,9 +404,8 @@ contract ZeusScriptTest is EOADeployer { function testZeusRequireMultisigEvent() public { vm.expectEmit(true, true, true, true); - emit ZeusRequireMultisig(address(0xabc), Operation.Call); - // Operation.Call and Operation.DelegateCall are from the enum defined in ZeusScript - emit ZeusRequireMultisig(address(0xabc), Operation.Call); + emit ZeusRequireMultisig(address(0xabc), Encode.Operation.Call); + emit ZeusRequireMultisig(address(0xabc), Encode.Operation.Call); } function testZeusDeployEvent() public { @@ -519,8 +518,8 @@ contract ZeusScriptTest is EOADeployer { function testZeusRequireMultisigEventDelegateCall() public { // Emit with DelegateCall to cover enum branch vm.expectEmit(true, true, true, true); - emit ZeusRequireMultisig(address(0xabc), Operation.DelegateCall); - emit ZeusRequireMultisig(address(0xabc), Operation.DelegateCall); + emit ZeusRequireMultisig(address(0xabc), Encode.Operation.DelegateCall); + emit ZeusRequireMultisig(address(0xabc), Encode.Operation.DelegateCall); } function testZeusDeployEventFalseSingleton() public {