From 65ed6481bfeed8fe6fa1b7e9a11de704fcfa1ac5 Mon Sep 17 00:00:00 2001
From: Josh Levine <joshuapark95@gmail.com>
Date: Wed, 7 Aug 2024 17:27:54 -0700
Subject: [PATCH] refactor: block data < 4 || > 0 and add tests

---
 script/RumpelWalletFactory.s.sol | 10 ++++-
 src/RumpelGuard.sol              |  6 +++
 test/RumpelWallet.t.sol          | 73 ++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/script/RumpelWalletFactory.s.sol b/script/RumpelWalletFactory.s.sol
index 8af458c..e827241 100644
--- 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
index c57a06a..44155f4 100644
--- 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
index 15bf427..edfb6ae 100644
--- 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 {}
 }