🔥 Changed static/call in safe7579 to use direct call to handler

rhinestonewtf · Mar 14, 2024 · 5c10bc7 · 5c10bc7
1 parent 0774d2b
commit 5c10bc7
Show file tree

Hide file tree

Showing 5 changed files with 92 additions and 71 deletions.
diff --git a/accounts/safe7579/README.md b/accounts/safe7579/README.md
diff --git a/accounts/safe7579/notes.txt b/accounts/safe7579/notes.txt
diff --git a/accounts/safe7579/src/core/ModuleManager.sol b/accounts/safe7579/src/core/ModuleManager.sol
@@ -107,6 +107,9 @@ abstract contract ModuleManager is AccessControl, Receiver, ExecutionHelper {
         isInstalled = $validators.contains({ account: msg.sender, entry: validator });
     }
 
+    /**
+     * Get paginated list of installed validators
+     */
     function getValidatorPaginated(
         address start,
         uint256 pageSize
@@ -240,10 +243,71 @@ abstract contract ModuleManager is AccessControl, Receiver, ExecutionHelper {
         if (handler == address(0)) revert NoFallbackHandler(msg.sig);
 
         if (calltype == CALLTYPE_STATIC) {
-            return _executeStaticReturnData(msg.sender, handler, 0, callData);
+            assembly {
+                // When compiled with the optimizer, the compiler relies on a certain assumptions on
+                // how
+                // the
+                // memory is used, therefore we need to guarantee memory safety (keeping the free
+                // memory
+                // point 0x40 slot intact,
+                // not going beyond the scratch space, etc)
+                // Solidity docs:
+                // https://docs.soliditylang.org/en/latest/assembly.html#memory-safety
+                function allocate(length) -> pos {
+                    pos := mload(0x40)
+                    mstore(0x40, add(pos, length))
+                }
+
+                let calldataPtr := allocate(calldatasize())
+                calldatacopy(calldataPtr, 0, calldatasize())
+
+                // The msg.sender address is shifted to the left by 12 bytes to remove the padding
+                // Then the address without padding is stored right after the calldata
+                let senderPtr := allocate(20)
+                mstore(senderPtr, shl(96, caller()))
+
+                // Add 20 bytes for the address appended add the end
+                let success :=
+                    staticcall(gas(), handler,  calldataPtr, add(calldatasize(), 20), 0, 0)
+
+                let returnDataPtr := allocate(returndatasize())
+                returndatacopy(returnDataPtr, 0, returndatasize())
+                if iszero(success) { revert(returnDataPtr, returndatasize()) }
+                return(returnDataPtr, returndatasize())
+            }
         }
         if (calltype == CALLTYPE_SINGLE) {
-            return _executeReturnData(msg.sender, handler, 0, callData);
+            assembly {
+                // When compiled with the optimizer, the compiler relies on a certain assumptions on
+                // how
+                // the
+                // memory is used, therefore we need to guarantee memory safety (keeping the free
+                // memory
+                // point 0x40 slot intact,
+                // not going beyond the scratch space, etc)
+                // Solidity docs:
+                // https://docs.soliditylang.org/en/latest/assembly.html#memory-safety
+                function allocate(length) -> pos {
+                    pos := mload(0x40)
+                    mstore(0x40, add(pos, length))
+                }
+
+                let calldataPtr := allocate(calldatasize())
+                calldatacopy(calldataPtr, 0, calldatasize())
+
+                // The msg.sender address is shifted to the left by 12 bytes to remove the padding
+                // Then the address without padding is stored right after the calldata
+                let senderPtr := allocate(20)
+                mstore(senderPtr, shl(96, caller()))
+
+                // Add 20 bytes for the address appended add the end
+                let success := call(gas(), handler, 0, calldataPtr, add(calldatasize(), 20), 0, 0)
+
+                let returnDataPtr := allocate(returndatasize())
+                returndatacopy(returnDataPtr, 0, returndatasize())
+                if iszero(success) { revert(returnDataPtr, returndatasize()) }
+                return(returnDataPtr, returndatasize())
+            }
         }
 
         if (calltype == CALLTYPE_DELEGATECALL) {

diff --git a/accounts/safe7579/test/SafeERC7579.t.sol b/accounts/safe7579/test/SafeERC7579.t.sol
@@ -7,6 +7,7 @@ import "erc7579/lib/ExecutionLib.sol";
 import { TestBaseUtil, MockTarget, MockFallback } from "./Base.t.sol";
 
 import "forge-std/console2.sol";
+CallType constant CALLTYPE_STATIC = CallType.wrap(0xFE);
 
 contract Safe7579Test is TestBaseUtil {
     MockTarget target;
@@ -154,10 +155,26 @@ contract Safe7579Test is TestBaseUtil {
         IERC7579Account(address(safe)).installModule(
             3, address(_fallback), abi.encode(MockFallback.target.selector, CALLTYPE_SINGLE, "")
         );
-        (uint256 ret, address msgSender) = MockFallback(address(safe)).target(1337);
+        (uint256 ret, address msgSender, address context) = MockFallback(address(safe)).target(1337);
 
         assertEq(ret, 1337);
-        assertEq(msgSender, address(safe));
+        assertEq(msgSender, address(safe7579));
+        assertEq(context, address(safe));
+
+        vm.prank(address(safe));
+        IERC7579Account(address(safe)).uninstallModule(
+            3, address(_fallback), abi.encode(MockFallback.target.selector, CALLTYPE_SINGLE, "")
+        );
+        vm.prank(address(safe));
+        IERC7579Account(address(safe)).installModule(
+            3, address(_fallback), abi.encode(MockFallback.target.selector, CALLTYPE_STATIC, "")
+        );
+        ( ret,  msgSender,  context) = MockFallback(address(safe)).target(1337);
+        assertEq(ret, 1337);
+        assertEq(msgSender, address(safe7579));
+        assertEq(context, address(safe));
+
+
 
         vm.prank(address(safe));
         IERC7579Account(address(safe)).installModule(

diff --git a/accounts/safe7579/test/mocks/MockFallback.sol b/accounts/safe7579/test/mocks/MockFallback.sol
@@ -4,10 +4,16 @@ pragma solidity ^0.8.23;
 import { HandlerContext } from "@safe-global/safe-contracts/contracts/handler/HandlerContext.sol";
 import { MockFallback as MockFallbackBase } from "@rhinestone/modulekit/src/mocks/MockFallback.sol";
 
+import "forge-std/console2.sol";
+
 contract MockFallback is MockFallbackBase, HandlerContext {
-    function target(uint256 value) external returns (uint256 _value, address msgSender) {
+    function target(uint256 value)
+        external
+        returns (uint256 _value, address msgSender, address msgSenderContext)
+    {
         _value = value;
         msgSender = msg.sender;
+        msgSenderContext = _msgSender();
     }
 
     function target2(uint256 value)