Merge pull request #1 from sense-finance/dev

New Features for 07/22 Audit

script/RumpelWalletFactory.s.sol

@@ -15,19 +15,23 @@ contract RumpelWalletFactoryScripts is Script {
     address public MAINNET_RUMPEL_VAULT = 0x1EeEBa76f211C4Dce994b9c5A74BDF25DB649Fa1;
     address public MAINNET_SAFE_SINGLETON = 0xd9Db270c1B5E3Bd161E8c8503c55cEABeE709552;
     ISafeProxyFactory public MAINNET_SAFE_PROXY_FACTORY = ISafeProxyFactory(0xa6B71E26C5e0845f74c812102Ca7114b6a896AB2);
+    address public MAINNET_COMPATIBILITY_FALLBACK = 0xf48f2B2d2a534e402487b3ee7C18c33Aec0Fe5e4;
+
+    address public MAINNET_SIGN_MESSAGE_LIB = 0xA65387F16B013cf2Af4605Ad8aA5ec25a2cbA3a2;
 
     function setUp() public {}
 
     function run(address admin) public returns (RumpelModule, RumpelGuard, RumpelWalletFactory) {
-        RumpelModule rumpelModule = new RumpelModule();
+        RumpelModule rumpelModule = new RumpelModule(MAINNET_SIGN_MESSAGE_LIB);
 
-        RumpelGuard rumpelGuard = new RumpelGuard();
+        RumpelGuard rumpelGuard = new RumpelGuard(MAINNET_SIGN_MESSAGE_LIB);
 
         // Script that will be delegatecalled to enable modules
         InitializationScript initializationScript = new InitializationScript();
 
         RumpelWalletFactory rumpelWalletFactory = new RumpelWalletFactory(
             MAINNET_SAFE_PROXY_FACTORY,
+            MAINNET_COMPATIBILITY_FALLBACK,
             MAINNET_SAFE_SINGLETON,
             address(rumpelModule),
             address(rumpelGuard),

src/InitializationScript.sol

@@ -2,12 +2,44 @@
 pragma solidity ^0.8.19;
 
 import {ISafe} from "./interfaces/external/ISafe.sol";
+import {Enum} from "./interfaces/external/ISafe.sol";
+import {RumpelGuard} from "./RumpelGuard.sol";
 
 /// @notice Delegatecall script to initialize Safes.
 contract InitializationScript {
+    error InitializationFailed();
+
+    event InitialCall(address to, bytes data);
+
+    struct InitCall {
+        address to;
+        bytes data;
+    }
+
     /// @notice This function is called via delegatecall by newly created Safes.
-    function initialize(address module, address guard) external {
-        ISafe(address(this)).enableModule(module);
-        ISafe(address(this)).setGuard(guard);
+    function initialize(address module, address guard, InitCall[] memory initCalls) external {
+        ISafe safe = ISafe(address(this));
+        safe.enableModule(module);
+        safe.setGuard(guard);
+
+        // Arbitrary initial calls.
+        for (uint256 i = 0; i < initCalls.length; i++) {
+            address to = initCalls[i].to;
+            bytes memory data = initCalls[i].data;
+
+            // Check each tx with the guard.
+            RumpelGuard(guard).checkTransaction(
+                to, 0, data, Enum.Operation.Call, 0, 0, 0, address(0), payable(address(0)), bytes(""), address(0)
+            );
+
+            bool success;
+            assembly {
+                success := call(sub(gas(), 500), to, 0, add(data, 0x20), mload(data), 0, 0)
+            }
+
+            if (!success) revert InitializationFailed();
+
+            emit InitialCall(to, data);
+        }
     }
 }

src/RumpelGuard.sol

@@ -11,6 +11,8 @@ import {IGuard} from "./interfaces/external/IGuard.sol";
 contract RumpelGuard is Ownable, IGuard {
     mapping(address => mapping(bytes4 => AllowListState)) public allowedCalls; // target => functionSelector => allowListState
 
+    address public immutable signMessageLib;
+
     enum AllowListState {
         OFF,
         ON,
@@ -22,15 +24,17 @@ contract RumpelGuard is Ownable, IGuard {
     error CallNotAllowed(address target, bytes4 functionSelector);
     error PermanentlyOn();
 
-    constructor() Ownable(msg.sender) {}
+    constructor(address _signMessageLib) Ownable(msg.sender) {
+        signMessageLib = _signMessageLib;
+    }
 
     /// @notice Called by the Safe contract before a transaction is executed.
     /// @dev Safe user execution hook that blocks all calls by default, including delegatecalls, unless explicitly added to the allowlist.
     function checkTransaction(
         address to,
         uint256,
         bytes memory data,
-        Enum.Operation,
+        Enum.Operation operation,
         uint256,
         uint256,
         uint256,
@@ -41,6 +45,15 @@ contract RumpelGuard is Ownable, IGuard {
     ) external view {
         bytes4 functionSelector = bytes4(data);
 
+        // Only allow delegatecalls to the signMessageLib.
+        if (operation == Enum.Operation.DelegateCall) {
+            if (to == signMessageLib) {
+                return;
+            } else {
+                revert CallNotAllowed(to, functionSelector);
+            }
+        }
+
         if (allowedCalls[to][functionSelector] == AllowListState.OFF) {
             revert CallNotAllowed(to, functionSelector);
         }
@@ -59,7 +72,7 @@ contract RumpelGuard is Ownable, IGuard {
     /// @notice Enable or disable Safes from calling a function.
     /// @dev Scoped to <address>.<selector>, so all calls to added address <> selector pairs are allowed.
     /// @dev Function arguments aren't checked, so any arguments are allowed for the enabled functions.
-    /// @dev Calls can be enabled, disabled, or permanently enabled, that last of which guarntees the call can't be rugged.
+    /// @dev Calls can be enabled, disabled, or permanently enabled, that last of which guarantees the call can't be rugged.
     function setCallAllowed(address target, bytes4 functionSelector, AllowListState allowListState)
         external
         onlyOwner

src/RumpelModule.sol

@@ -9,12 +9,14 @@ import {ISafe} from "./interfaces/external/ISafe.sol";
 /// @notice Rumpel Safe Module allowing an admin to execute calls on behalf of the Safe.
 contract RumpelModule is Ownable {
     mapping(address => mapping(bytes4 => bool)) public blockedModuleCalls; // target => functionSelector => blocked
+    address public immutable signMessageLib;
 
     struct Call {
         ISafe safe;
         address to;
         uint256 value;
         bytes data;
+        Enum.Operation operation;
     }
 
     event ExecutionFromModule(ISafe indexed safe, address indexed target, uint256 value, bytes data);
@@ -23,7 +25,9 @@ contract RumpelModule is Ownable {
     error ExecFailed(ISafe safe, address target, bytes data);
     error CallBlocked(address target, bytes4 functionSelector);
 
-    constructor() Ownable(msg.sender) {}
+    constructor(address _signMessageLib) Ownable(msg.sender) {
+        signMessageLib = _signMessageLib;
+    }
 
     /// @notice Execute a series of calls through Safe contracts.
     /// @param calls An array of Call structs containing the details of each call to be executed.
@@ -38,7 +42,14 @@ contract RumpelModule is Ownable {
                 revert CallBlocked(call.to, bytes4(call.data));
             }
 
-            bool success = call.safe.execTransactionFromModule(call.to, call.value, call.data, Enum.Operation.Call); // No delegatecalls
+            // Only allow delegatecalls to the signMessageLib.
+            if (call.operation == Enum.Operation.DelegateCall) {
+                if (call.to != signMessageLib) {
+                    revert CallBlocked(call.to, bytes4(call.data));
+                }
+            }
+
+            bool success = call.safe.execTransactionFromModule(call.to, call.value, call.data, call.operation);
 
             if (!success) {
                 revert ExecFailed(call.safe, call.to, call.data);
@@ -52,7 +63,7 @@ contract RumpelModule is Ownable {
         }
     }
 
-    /// @notice Prevent call from being executed via the module permanently. Useful as an assurance that e.g. and admin will never transfer a user's USDC.
+    /// @notice Prevent call from being executed via the module permanently. Useful as an assurance that e.g. an admin will never transfer a user's USDC.
     /// @dev Scoped to <address>.<selector>, so all calls to added address <> selector pairs are blocked.
     /// @dev To block calls to arbitrary Safes, to prevent an admin from updating config, the Zero address is used for the target.
     function addBlockedModuleCall(address target, bytes4 functionSelector) external onlyOwner {

src/RumpelWalletFactory.sol

@@ -9,9 +9,10 @@ import {InitializationScript} from "./InitializationScript.sol";
 
 /// @notice Factory to create Rumpel Wallets; Safes with the Rumpel Guard and Rumpel Module added on.
 contract RumpelWalletFactory is Ownable, Pausable {
-    uint256 public saltNonce;
+    mapping(address => uint256) public saltNonce;
 
     ISafeProxyFactory public proxyFactory;
+    address public compatibilityFallback;
     address public safeSingleton;
     address public rumpelModule;
     address public rumpelGuard;
@@ -24,34 +25,40 @@ contract RumpelWalletFactory is Ownable, Pausable {
 
     constructor(
         ISafeProxyFactory _proxyFactory,
+        address _compatibilityFallback,
         address _safeSingleton,
         address _rumpelModule,
         address _rumpelGuard,
         address _initializationScript
     ) Ownable(msg.sender) {
         proxyFactory = _proxyFactory;
+        compatibilityFallback = _compatibilityFallback;
         safeSingleton = _safeSingleton;
         rumpelModule = _rumpelModule;
         rumpelGuard = _rumpelGuard;
         initializationScript = _initializationScript;
     }
 
     /// @notice Create a Safe with the Rumpel Module and Rumpel Guard added.
-    function createWallet(address[] calldata owners, uint256 threshold) external whenNotPaused returns (address) {
+    function createWallet(
+        address[] calldata owners,
+        uint256 threshold,
+        InitializationScript.InitCall[] calldata initCalls
+    ) external whenNotPaused returns (address) {
         address safe = proxyFactory.createProxyWithNonce(
             safeSingleton,
             abi.encodeWithSelector(
                 ISafe.setup.selector,
                 owners,
                 threshold,
                 initializationScript, // Contract with initialization logic
-                abi.encodeWithSelector(InitializationScript.initialize.selector, rumpelModule, rumpelGuard), // Initializing call to add module and guard
-                address(0), // fallbackHandler
+                abi.encodeWithSelector(InitializationScript.initialize.selector, rumpelModule, rumpelGuard, initCalls), // Add module and guard + initial calls
+                compatibilityFallback, // fallbackHandler
                 address(0), // paymentToken
                 0, // payment
                 address(0) // paymentReceiver
             ),
-            saltNonce++
+            saltNonce[msg.sender]++ // For deterministic address generation
         );
 
         emit SafeCreated(safe, owners, threshold);

src/interfaces/external/ISafe.sol

@@ -63,4 +63,18 @@ interface ISafe {
     function isModuleEnabled(address module) external view returns (bool);
 
     function nonce() external view returns (uint256);
+
+    function execute(address to, uint256 value, bytes memory data, Enum.Operation operation, uint256 txGas)
+        external
+        returns (bool success);
+
+    function domainSeparator() external view returns (bytes32);
+
+    function signedMessages(bytes32 messageHash) external returns (uint256);
+
+    // Fallback handler functions
+
+    function isValidSignature(bytes32 _dataHash, bytes calldata _signature) external view returns (bytes4);
+
+    function isValidSignature(bytes calldata _data, bytes calldata _signature) external view returns (bytes4);
 }

src/interfaces/external/ISafeProxyFactory.sol

@@ -5,4 +5,5 @@ interface ISafeProxyFactory {
     function createProxyWithNonce(address _singleton, bytes memory _initializer, uint256 _saltNonce)
         external
         returns (address);
+    function proxyCreationCode() external pure returns (bytes memory);
 }

src/interfaces/external/ISignMessageLib.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+interface ISignMessageLib {
+    function signMessage(bytes calldata digest) external;
+}