fix: catch data in all cases (#151)

erc6900 · Aug 22, 2024 · 0c18f92 · 0c18f92
1 parent 78bc489
commit 0c18f92
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 6 deletions.
diff --git a/src/account/ModuleManagerInternals.sol b/src/account/ModuleManagerInternals.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.25;
 import {ERC165Checker} from "@openzeppelin/contracts/utils/introspection/ERC165Checker.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 
+import {collectReturnData} from "../helpers/CollectReturnData.sol";
 import {MAX_PRE_VALIDATION_HOOKS} from "../helpers/Constants.sol";
 import {HookConfigLib} from "../helpers/HookConfigLib.sol";
 import {KnownSelectors} from "../helpers/KnownSelectors.sol";
@@ -170,7 +171,8 @@ abstract contract ModuleManagerInternals is IModularAccount {
         // Initialize the module storage for the account.
         // solhint-disable-next-line no-empty-blocks
         try IModule(module).onInstall(moduleInstallData) {}
-        catch (bytes memory revertReason) {
+        catch {
+            bytes memory revertReason = collectReturnData();
             revert ModuleInstallCallbackFailed(module, revertReason);
         }
 

diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
@@ -11,6 +11,7 @@ import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeab
 import {IERC165} from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 
+import {collectReturnData} from "../helpers/CollectReturnData.sol";
 import {DIRECT_CALL_VALIDATION_ENTITYID} from "../helpers/Constants.sol";
 import {HookConfig, HookConfigLib} from "../helpers/HookConfigLib.sol";
 import {ModuleEntityLib} from "../helpers/ModuleEntityLib.sol";
@@ -460,8 +461,8 @@ contract UpgradeableModularAccount is
             bytes memory returnData
         ) {
             preExecHookReturnData = returnData;
-        } catch (bytes memory revertReason) {
-            // TODO: same issue with EP0.6 - we can't do bytes4 error codes in modules
+        } catch {
+            bytes memory revertReason = collectReturnData();
             revert PreExecHookReverted(module, entityId, revertReason);
         }
     }
@@ -483,7 +484,8 @@ contract UpgradeableModularAccount is
             (address module, uint32 entityId) = postHookToRun.postExecHook.unpack();
             // solhint-disable-next-line no-empty-blocks
             try IExecutionHookModule(module).postExecutionHook(entityId, postHookToRun.preExecHookReturnData) {}
-            catch (bytes memory revertReason) {
+            catch {
+                bytes memory revertReason = collectReturnData();
                 revert PostExecHookReverted(module, entityId, revertReason);
             }
         }
@@ -500,8 +502,9 @@ contract UpgradeableModularAccount is
         )
         // forgefmt: disable-start
         // solhint-disable-next-line no-empty-blocks
-        {} catch (bytes memory revertReason){
+        {} catch{
         // forgefmt: disable-end
+            bytes memory revertReason = collectReturnData();
             revert PreRuntimeValidationHookFailed(hookModule, hookEntityId, revertReason);
         }
     }
@@ -585,8 +588,9 @@ contract UpgradeableModularAccount is
         )
         // forgefmt: disable-start
         // solhint-disable-next-line no-empty-blocks
-        {} catch (bytes memory revertReason){
+        {} catch{
         // forgefmt: disable-end
+            bytes memory revertReason = collectReturnData();
             revert RuntimeValidationFunctionReverted(module, entityId, revertReason);
         }
     }

diff --git a/src/helpers/CollectReturnData.sol b/src/helpers/CollectReturnData.sol
@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.25;
+
+function collectReturnData() pure returns (bytes memory returnData) {
+    assembly ("memory-safe") {
+        // Allocate a buffer of that size, advancing the memory pointer to the nearest word
+        returnData := mload(0x40)
+        mstore(returnData, returndatasize())
+        mstore(0x40, and(add(add(returnData, returndatasize()), 0x3f), not(0x1f)))
+
+        // Copy over the return data
+        returndatacopy(add(returnData, 0x20), 0, returndatasize())
+    }
+}