refactor: block data < 4 || > 0 and add tests

sense-finance · Aug 8, 2024 · 65ed648 · 65ed648
1 parent 953b264
commit 65ed648
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 1 deletion.
diff --git a/script/RumpelWalletFactory.s.sol b/script/RumpelWalletFactory.s.sol
@@ -22,6 +22,14 @@ contract RumpelWalletFactoryScripts is Script {
 
     function setUp() public {}
 
+    function run() public {
+        uint256 deployerPrivateKey = vm.envUint("DEPLOYER_PRIVATE_KEY");
+
+        vm.startBroadcast(deployerPrivateKey);
+        run(msg.sender);
+        vm.stopBroadcast();
+    }
+
     function run(address admin) public returns (RumpelModule, RumpelGuard, RumpelWalletFactory) {
         RumpelModule rumpelModule = new RumpelModule(MAINNET_SIGN_MESSAGE_LIB);
 
@@ -42,7 +50,7 @@ contract RumpelWalletFactoryScripts is Script {
         // Prevent us from just swapping out the Rumpel Module for one without a blocklist.
         rumpelModule.addBlockedModuleCall(address(0), ISafe.enableModule.selector);
 
-        // Allow safes to delegate pToken claiming
+        // Allow Safes to delegate pToken claiming
         rumpelGuard.setCallAllowed(
             MAINNET_RUMPEL_VAULT, PointTokenVault.trustClaimer.selector, RumpelGuard.AllowListState.ON
         );

diff --git a/src/RumpelGuard.sol b/src/RumpelGuard.sol
@@ -43,6 +43,12 @@ contract RumpelGuard is Ownable, IGuard {
         bytes memory,
         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.
+        if (data.length > 0 && data.length < 4) {
+            revert CallNotAllowed(to, bytes4(data));
+        }
+
         bytes4 functionSelector = bytes4(data);
 
         // Only allow delegatecalls to the signMessageLib.

diff --git a/test/RumpelWallet.t.sol b/test/RumpelWallet.t.sol
@@ -386,6 +386,77 @@ contract RumpelWalletTest is Test {
         );
     }
 
+    function test_RumpelWalletDisallowSmallData() 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));
+
+        // 3 bytes. Will be padded out with two 0s when cast to bytes4
+        bytes memory smallData = new bytes(3);
+        smallData[0] = bytes1(uint8(1));
+        smallData[1] = bytes1(uint8(2));
+        smallData[2] = bytes1(uint8(3));
+
+        vm.prank(admin);
+        rumpelGuard.setCallAllowed(address(counter), bytes4(smallData), RumpelGuard.AllowListState.ON);
+
+        // Will revert even though the data has been allowed, because the data is too small
+        vm.expectRevert(
+            abi.encodeWithSelector(RumpelGuard.CallNotAllowed.selector, address(counter), bytes4(smallData))
+        );
+        this._execSafeTx(safe, address(counter), 0, smallData, Enum.Operation.Call);
+    }
+
+    function test_RumpelWalletAllowETHTransfers() 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 zeroData = new bytes(0);
+
+        // Enable ETH transfers
+        vm.prank(admin);
+        rumpelGuard.setCallAllowed(address(counter), bytes4(zeroData), 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);
+
+        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);
+    }
+
+    function test_RumpelWalletAllowZeroData() 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));
+
+        // 3 bytes. Will be padded out with two 0s when cast to bytes4
+        bytes memory smallData = new bytes(3);
+        smallData[0] = bytes1(uint8(1));
+        smallData[1] = bytes1(uint8(2));
+        smallData[2] = bytes1(uint8(3));
+
+        vm.prank(admin);
+        rumpelGuard.setCallAllowed(address(counter), bytes4(smallData), RumpelGuard.AllowListState.ON);
+
+        // Will revert even though the data has been allowed, because the data is too small
+        vm.expectRevert(
+            abi.encodeWithSelector(RumpelGuard.CallNotAllowed.selector, address(counter), bytes4(smallData))
+        );
+        this._execSafeTx(safe, address(counter), 0, smallData, Enum.Operation.Call);
+    }
+
     function test_GuardPermanentlyAllowedCall() public {
         address[] memory owners = new address[](1);
         owners[0] = address(alice);
@@ -731,4 +802,6 @@ contract Counter {
     function fail() external {
         revert("fail");
     }
+
+    fallback() external payable {}
 }