feat: add guard levels for safe self config updates and eth transfers…

… broadly
sense-finance · Aug 24, 2024 · b3fb665 · b3fb665
1 parent d0f9a14
commit b3fb665
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 4 deletions.
diff --git a/src/RumpelGuard.sol b/src/RumpelGuard.sol
@@ -5,6 +5,7 @@ import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 
 import {Enum} from "./interfaces/external/ISafe.sol";
 import {IGuard} from "./interfaces/external/IGuard.sol";
+import {console} from "forge-std/console.sol";
 
 /// @notice Rumpel Safe Guard with a blocklist for the Rumpel Wallet.
 /// @dev Compatible with Safe v1.3.0-libs.0, the last Safe Ethereum mainnet release, so it can't use module execution hooks.
@@ -44,7 +45,7 @@ contract RumpelGuard is Ownable, IGuard {
         address
     ) external view {
         // Disallow calls with function selectors that will be padded with 0s.
-        // Allow calls to data with length 0, so that the call can be used for reentrancy.
+        // Allow calls with data length 0 for ETH transfers.
         if (data.length > 0 && data.length < 4) {
             revert CallNotAllowed(to, bytes4(data));
         }
@@ -60,8 +61,23 @@ contract RumpelGuard is Ownable, IGuard {
             }
         }
 
-        if (allowedCalls[to][functionSelector] == AllowListState.OFF) {
-            revert CallNotAllowed(to, functionSelector);
+        bool toSafe = msg.sender == to;
+
+        if (toSafe) {
+            // If this transaction is to a Safe itself, to e.g. update config, we check the zero address for allowed calls.
+            if (allowedCalls[address(0)][functionSelector] == AllowListState.OFF) {
+                revert CallNotAllowed(to, functionSelector);
+            }
+        } else if (data.length == 0) {
+            // If this transaction is a simple ETH transfer, we check the zero address with the zero function selector to see if it's allowed.
+            if (allowedCalls[address(0)][bytes4(0)] == AllowListState.OFF) {
+                revert CallNotAllowed(address(0), bytes4(0));
+            }
+        } else {
+            // For all other calls, we check the allowedCalls mapping normally.
+            if (allowedCalls[to][functionSelector] == AllowListState.OFF) {
+                revert CallNotAllowed(to, functionSelector);
+            }
         }
     }
 

diff --git a/src/interfaces/external/ISafe.sol b/src/interfaces/external/ISafe.sol
@@ -74,6 +74,8 @@ interface ISafe {
 
     function signedMessages(bytes32 messageHash) external returns (uint256);
 
+    function addOwnerWithThreshold(address owner, uint256 _threshold) external;
+
     // Fallback handler functions
 
     function isValidSignature(bytes32 _dataHash, bytes calldata _signature) external view returns (bytes4);

diff --git a/test/RumpelWallet.t.sol b/test/RumpelWallet.t.sol
@@ -409,18 +409,47 @@ contract RumpelWalletTest is Test {
 
         // Enable ETH transfers
         vm.prank(admin);
-        rumpelGuard.setCallAllowed(address(counter), bytes4(zeroData), RumpelGuard.AllowListState.ON);
+        rumpelGuard.setCallAllowed(address(0), bytes4(0), RumpelGuard.AllowListState.ON);
 
         // Mint 1 ETH to the safe
         vm.deal(address(safe), 1 ether);
 
         assertEq(address(safe).balance, 1 ether);
         assertEq(address(counter).balance, 0);
 
+        // Transfer to contract
         this._execSafeTx(safe, address(counter), 0.1 ether, zeroData, Enum.Operation.Call);
 
         assertEq(address(safe).balance, 0.9 ether);
         assertEq(address(counter).balance, 0.1 ether);
+
+        // Transfer to address
+        this._execSafeTx(safe, address(alice), 0.1 ether, zeroData, Enum.Operation.Call);
+
+        assertEq(address(safe).balance, 0.8 ether);
+        assertEq(address(alice).balance, 0.1 ether);
+    }
+
+    function test_RumpelWalletConfigUpdateAuth() public {
+        address[] memory owners = new address[](1);
+        owners[0] = address(alice);
+
+        InitializationScript.InitCall[] memory initCalls = new InitializationScript.InitCall[](0);
+        ISafe safe = ISafe(rumpelWalletFactory.createWallet(owners, 1, initCalls));
+
+        bytes memory addOwnerData = abi.encodeCall(ISafe.addOwnerWithThreshold, (makeAddr("bob"), 1));
+
+        // Try to add an owner to the safe wallet
+        vm.expectRevert(
+            abi.encodeWithSelector(RumpelGuard.CallNotAllowed.selector, address(safe), bytes4(addOwnerData))
+        );
+        this._execSafeTx(safe, address(safe), 0, addOwnerData, Enum.Operation.Call);
+
+        vm.prank(admin);
+        rumpelGuard.setCallAllowed(address(0), bytes4(addOwnerData), RumpelGuard.AllowListState.ON);
+        this._execSafeTx(safe, address(safe), 0, addOwnerData, Enum.Operation.Call);
+
+        assertEq(safe.getOwners().length, 2);
     }
 
     function test_RumpelWalletAllowZeroData() public {