Skip to content

Commit

Permalink
feat: skip return data on self call
Browse files Browse the repository at this point in the history
  • Loading branch information
adamegyed committed Oct 11, 2024
1 parent e32e33f commit 8a95166
Show file tree
Hide file tree
Showing 28 changed files with 54 additions and 63 deletions.
Original file line number Diff line number Diff line change
@@ -1 +1 @@
93449
93127
2 changes: 1 addition & 1 deletion .forge-snapshots/ModularAccount_Runtime_Erc20Transfer.snap
Original file line number Diff line number Diff line change
@@ -1 +1 @@
78730
78552
Original file line number Diff line number Diff line change
@@ -1 +1 @@
424033
424043
Original file line number Diff line number Diff line change
@@ -1 +1 @@
54870
54701
Original file line number Diff line number Diff line change
@@ -1 +1 @@
79712
79540
Original file line number Diff line number Diff line change
@@ -1 +1 @@
113276
113098
2 changes: 1 addition & 1 deletion .forge-snapshots/ModularAccount_UserOp_BatchTransfers.snap
Original file line number Diff line number Diff line change
@@ -1 +1 @@
200346
198425
2 changes: 1 addition & 1 deletion .forge-snapshots/ModularAccount_UserOp_Erc20Transfer.snap
Original file line number Diff line number Diff line change
@@ -1 +1 @@
185789
185394
2 changes: 1 addition & 1 deletion .forge-snapshots/ModularAccount_UserOp_NativeTransfer.snap
Original file line number Diff line number Diff line change
@@ -1 +1 @@
161972
161675
Original file line number Diff line number Diff line change
@@ -1 +1 @@
196094
195915
Original file line number Diff line number Diff line change
@@ -1 +1 @@
227224
227039
Original file line number Diff line number Diff line change
@@ -1 +1 @@
257304
257016
Original file line number Diff line number Diff line change
@@ -1 +1 @@
89828
89614
Original file line number Diff line number Diff line change
@@ -1 +1 @@
75176
75094
Original file line number Diff line number Diff line change
@@ -1 +1 @@
422963
423069
Original file line number Diff line number Diff line change
@@ -1 +1 @@
51317
51244
Original file line number Diff line number Diff line change
@@ -1 +1 @@
80150
80038
Original file line number Diff line number Diff line change
@@ -1 +1 @@
113714
113596
Original file line number Diff line number Diff line change
@@ -1 +1 @@
195544
193755
Original file line number Diff line number Diff line change
@@ -1 +1 @@
181114
180839
Original file line number Diff line number Diff line change
@@ -1 +1 @@
529814
529934
Original file line number Diff line number Diff line change
@@ -1 +1 @@
157291
157114
Original file line number Diff line number Diff line change
@@ -1 +1 @@
196335
196336
Original file line number Diff line number Diff line change
@@ -1 +1 @@
227477
227472
Original file line number Diff line number Diff line change
@@ -1 +1 @@
253744
253636
28 changes: 21 additions & 7 deletions src/account/ModularAccountBase.sol
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ abstract contract ModularAccountBase is
bytes memory callData = ExecutionLib.getExecuteUOCallData(callBuffer, userOp.callData);

// Manually call self, without collecting return data unless there's a revert.
ExecutionLib.callSelfBubbleOnRevert(callData);
ExecutionLib.callBubbleOnRevert(address(this), 0, callData);

ExecutionLib.doCachedPostHooks(postHookData);
}
Expand All @@ -250,7 +250,12 @@ abstract contract ModularAccountBase is
wrapNativeFunction
returns (bytes memory result)
{
result = ExecutionLib.exec(target, value, data);
ExecutionLib.callBubbleOnRevertTransient(target, value, data);

// Only return data if not called by the EntryPoint
if (msg.sender != address(_ENTRY_POINT)) {
result = ExecutionLib.collectReturnData();
}
}

/// @inheritdoc IModularAccount
Expand All @@ -262,11 +267,20 @@ abstract contract ModularAccountBase is
wrapNativeFunction
returns (bytes[] memory results)
{
bool needReturnData = (msg.sender != address(_ENTRY_POINT));

uint256 callsLength = calls.length;
results = new bytes[](callsLength);

if (needReturnData) {
results = new bytes[](callsLength);
}

for (uint256 i = 0; i < callsLength; ++i) {
results[i] = ExecutionLib.exec(calls[i].target, calls[i].value, calls[i].data);
ExecutionLib.callBubbleOnRevertTransient(calls[i].target, calls[i].value, calls[i].data);

if (needReturnData) {
results[i] = ExecutionLib.collectReturnData();
}
}
}

Expand Down Expand Up @@ -301,8 +315,8 @@ abstract contract ModularAccountBase is

// Execute the call, reusing the already-allocated RT call buffers, if it exists.
// In practice, this is cheaper than attempting to coalesce the (possibly two) buffers.
bytes memory returnData =
ExecutionLib.exec(address(this), 0 wei, ExecutionLib.getCallData(rtCallBuffer, data));
ExecutionLib.callBubbleOnRevert(address(this), 0 wei, ExecutionLib.getCallData(rtCallBuffer, data));
bytes memory returnData = ExecutionLib.collectReturnData();

ExecutionLib.doCachedPostHooks(postHookData);

Expand Down Expand Up @@ -460,7 +474,7 @@ abstract contract ModularAccountBase is
validationData = uint256(deadline) << 160;

// Call `installValidation` on the account.
ExecutionLib.callSelfBubbleOnRevertTransient(encodedData[38:]);
ExecutionLib.callBubbleOnRevertTransient(address(this), 0, encodedData[38:]);

// Update the UserOp signature to the remaining bytes.
userOpSignature = userOp.signature[33 + encodedDataLength + deferredInstallSigLength:];
Expand Down
37 changes: 7 additions & 30 deletions src/libraries/ExecutionLib.sol
Original file line number Diff line number Diff line change
Expand Up @@ -31,39 +31,17 @@ library ExecutionLib {
error PreRuntimeValidationHookFailed(address module, uint32 entityId, bytes revertReason);
error RuntimeValidationFunctionReverted(address module, uint32 entityId, bytes revertReason);

/// @param target The address of the contract to call.
/// @param value The value to send with the call.
/// @param data The call data.
/// @return result The return data of the call, or the error message from the call if call reverts.
function exec(address target, uint256 value, bytes memory data) internal returns (bytes memory result) {
// Manually call, collecting return data.
assembly ("memory-safe") {
let success := call(gas(), target, value, add(data, 0x20), mload(data), codesize(), 0)

// Allocate space for the return data, advancing the memory pointer to the nearest word
result := mload(0x40)
mstore(0x40, and(add(add(result, returndatasize()), 0x3f), not(0x1f)))

// Copy the returned data to the allocated space.
mstore(result, returndatasize())
returndatacopy(add(result, 0x20), 0, returndatasize())

// Revert if the call was not successful.
if iszero(success) { revert(add(result, 0x20), returndatasize()) }
}
}

// Call the following function to address(this), without capturing any return data.
// Perform the following call, without capturing any return data.
// If the call reverts, the revert message will be directly bubbled up.
function callSelfBubbleOnRevert(bytes memory callData) internal {
function callBubbleOnRevert(address target, uint256 value, bytes memory callData) internal {
// Manually call, without collecting return data unless there's a revert.
assembly ("memory-safe") {
let success :=
call(
gas(),
address(),
target,
/*value*/
0,
value,
/*argOffset*/
add(callData, 0x20),
/*argSize*/
Expand All @@ -86,9 +64,8 @@ library ExecutionLib {
}
}

// Just like `callSelfBubbleOnRevert`, but with a `bytes calldata` parameter. This will be transiently stored
// in memory past-the-end of the used memory.
function callSelfBubbleOnRevertTransient(bytes calldata callData) internal {
// Transiently copy the call data to a memory, and perform a self-call.
function callBubbleOnRevertTransient(address target, uint256 value, bytes calldata callData) internal {
bytes memory callToSelf;

assembly ("memory-safe") {
Expand All @@ -99,7 +76,7 @@ library ExecutionLib {
calldatacopy(add(callToSelf, 0x20), callData.offset, callData.length)
}

callSelfBubbleOnRevert(callToSelf);
callBubbleOnRevert(target, value, callToSelf);
// Memory is discarded afterwards
}

Expand Down
2 changes: 1 addition & 1 deletion test/account/DirectCallsFromModule.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ contract DirectCallsFromModuleTest is AccountTestBase {
{
bytes memory encodedCall = abi.encodeCall(DirectCallModule.directCall, ());

vm.prank(address(entryPoint));
vm.prank(address(account1));
bytes memory result = account1.execute(address(_module), 0, encodedCall);

assertTrue(_module.preHookRan());
Expand Down

0 comments on commit 8a95166

Please sign in to comment.