From f4c19ba1035090dc891f8856b934c491aab5c21d Mon Sep 17 00:00:00 2001
From: Josh Levine <joshuapark95@gmail.com>
Date: Mon, 15 Jul 2024 14:09:28 -0700
Subject: [PATCH] feat: guard against all delegatecalls

---
 src/RumpelGuard.sol     |  4 ++--
 test/RumpelWallet.t.sol | 50 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/src/RumpelGuard.sol b/src/RumpelGuard.sol
index 6a5e7e4..7400a16 100644
--- a/src/RumpelGuard.sol
+++ b/src/RumpelGuard.sol
@@ -30,7 +30,7 @@ contract RumpelGuard is Ownable, IGuard {
         address to,
         uint256,
         bytes memory data,
-        Enum.Operation,
+        Enum.Operation operation,
         uint256,
         uint256,
         uint256,
@@ -41,7 +41,7 @@ contract RumpelGuard is Ownable, IGuard {
     ) external view {
         bytes4 functionSelector = bytes4(data);
 
-        if (allowedCalls[to][functionSelector] == AllowListState.OFF) {
+        if (operation == Enum.Operation.DelegateCall || allowedCalls[to][functionSelector] == AllowListState.OFF) {
             revert CallNotAllowed(to, functionSelector);
         }
     }
diff --git a/test/RumpelWallet.t.sol b/test/RumpelWallet.t.sol
index dff8242..481aca4 100644
--- a/test/RumpelWallet.t.sol
+++ b/test/RumpelWallet.t.sol
@@ -184,6 +184,48 @@ contract RumpelWalletTest is Test {
         assertEq(counter.count(), 1);
     }
 
+    function test_guardDisallowDelegateCall() public {
+        DelegateCallTestScript delegateCallTestScript = new DelegateCallTestScript();
+
+        address[] memory owners = new address[](1);
+        owners[0] = address(alice);
+
+        ISafe safe = ISafe(rumpelWalletFactory.createWallet(owners, 1));
+
+        // Enable call to the delegate call script
+        vm.prank(admin);
+        rumpelGuard.setCallAllowed(
+            address(delegateCallTestScript), DelegateCallTestScript.echo.selector, RumpelGuard.AllowListState.ON
+        );
+
+        // Build a delegate call transaction
+        SafeTX memory safeTX = SafeTX({
+            to: address(delegateCallTestScript),
+            value: 0,
+            data: abi.encodeCall(DelegateCallTestScript.echo, (123)),
+            operation: Enum.Operation.DelegateCall
+        });
+
+        uint256 nonce = safe.nonce();
+
+        bytes32 txHash = safe.getTransactionHash(
+            safeTX.to, safeTX.value, safeTX.data, safeTX.operation, 0, 0, 0, address(0), payable(address(0)), nonce
+        );
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(alicePk, txHash);
+        bytes memory signature = abi.encodePacked(r, s, v);
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                RumpelGuard.CallNotAllowed.selector,
+                address(delegateCallTestScript),
+                bytes4(abi.encodeCall(DelegateCallTestScript.echo, (123)))
+            )
+        );
+        safe.execTransaction(
+            safeTX.to, safeTX.value, safeTX.data, safeTX.operation, 0, 0, 0, address(0), payable(address(0)), signature
+        );
+    }
+
     function test_guardPermanentlyAllowedCall() public {
         address[] memory owners = new address[](1);
         owners[0] = address(alice);
@@ -463,6 +505,14 @@ contract MockExternalProtocol {
     }
 }
 
+contract DelegateCallTestScript {
+    event Echo(uint256);
+
+    function echo(uint256 num) external {
+        emit Echo(num);
+    }
+}
+
 contract Counter {
     uint256 public count;