chore: add module guard test, refactor sweep

src/InitializationScript.sol

@@ -3,7 +3,7 @@ pragma solidity ^0.8.13;
 
 import {ISafe} from "./interfaces/external/ISafe.sol";
 
-// For delegate calls from the safe
+// For delegatecalls from the safe
 contract InitializationScript {
     function initialize(address module, address guard) public {
         ISafe(address(this)).enableModule(module);

src/RumpelGuard.sol

@@ -1,20 +1,18 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.13;
 
-import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
 
 import {Enum} from "./interfaces/external/ISafe.sol";
 import {IGuard} from "./interfaces/external/IGuard.sol";
 
-contract RumpelGuard is Ownable, IGuard {
+contract RumpelGuard is AccessControl, IGuard {
     error CallNotAllowed();
 
-    constructor() Ownable(msg.sender) {}
-
     mapping(address => mapping(bytes4 => bool)) public allowedTargets;
 
-    // TODO: should we allow more fine-grained control for which args are allowed?
-    function setCallAllowed(address target, bytes4 functionSig, bool allow) public onlyOwner {
+    // TODO: should we allow more fine-grained control on which args are allowed?
+    function setCallAllowed(address target, bytes4 functionSig, bool allow) public onlyRole(DEFAULT_ADMIN_ROLE) {
         allowedTargets[target][functionSig] = allow;
     }
 
@@ -38,8 +36,8 @@ contract RumpelGuard is Ownable, IGuard {
 
     function checkAfterExecution(bytes32, bool) external view {}
 
-    function supportsInterface(bytes4 interfaceId) external pure returns (bool) {
-        return interfaceId == type(IGuard).interfaceId;
+    function supportsInterface(bytes4 interfaceId) public view override returns (bool) {
+        return super.supportsInterface(interfaceId) || interfaceId == type(IGuard).interfaceId;
     }
 
     fallback() external {}

src/RumpelModule.sol

@@ -8,6 +8,8 @@ import {Enum} from "./interfaces/external/ISafe.sol";
 import {ISafe} from "./interfaces/external/ISafe.sol";
 
 contract RumpelModule is AccessControl {
+    error ExecFailed();
+
     bytes32 public constant SWEEPER_ROLE = keccak256("SWEEPER_ROLE");
     address public rumpelVault;
 
@@ -21,17 +23,25 @@ contract RumpelModule is AccessControl {
         public
         virtual
         onlyRole(DEFAULT_ADMIN_ROLE)
-        returns (bool)
     {
-        return safe.execTransactionFromModule(to, value, data, operation);
+        bool success = safe.execTransactionFromModule(to, value, data, operation);
+        if (!success) {
+            revert ExecFailed();
+        }
+    }
+
+    struct Sweep {
+        ISafe safe;
+        ERC20 token;
+        uint256 amount;
     }
 
-    function sweep(ISafe[] memory safes, ERC20 token) public virtual onlyRole(SWEEPER_ROLE) {
-        for (uint256 i = 0; i < safes.length; i++) {
-            safes[i].execTransactionFromModule(
-                address(token),
+    function sweep(Sweep[] memory sweeps) public virtual onlyRole(SWEEPER_ROLE) {
+        for (uint256 i = 0; i < sweeps.length; i++) {
+            sweeps[i].safe.execTransactionFromModule(
+                address(sweeps[i].token),
                 0,
-                abi.encodeCall(ERC20.transfer, (rumpelVault, token.balanceOf(address(safes[i])))),
+                abi.encodeCall(ERC20.transfer, (rumpelVault, sweeps[i].amount)),
                 Enum.Operation.Call
             );
         }

src/RumpelWallet.sol → src/RumpelWalletFactory.sol

@@ -8,9 +8,7 @@ import {ISafeProxyFactory} from "./interfaces/external/ISafeProxyFactory.sol";
 import {InitializationScript} from "./InitializationScript.sol";
 import {RumpelModule} from "./RumpelModule.sol";
 
-// delegate claiming fnc?
-// more admin controls on the guard?
-// * roles for different interactions from e.g. the sweeper contract
+// delegate claiming fnc in the vault?
 
 contract RumpelWalletFactory is Ownable {
     uint256 public saltNonce;
@@ -49,15 +47,15 @@ contract RumpelWalletFactory is Ownable {
             ISafe.setup.selector,
             owners,
             threshold,
-            initializationScript, // Target contract we delegatecall for initialization
+            initializationScript, // Target contract we delegatecall to for initialization
             abi.encodeWithSelector(InitializationScript.initialize.selector, rumpelModule, rumpelGuard), // Initializing call to enable module and guard
             fallbackHandler,
             paymentToken,
             payment,
             paymentReceiver
         );
 
-        return proxyFactory.createProxyWithNonce(singleton, initializer, saltNonce++);
+        return proxyFactory.createProxyWithNonce(singleton, initializer, saltNonce++); // TODO: do we want to limit each owner to one wallet via nonce?
     }
 
     // Admin ---

test/RumpelWallet.t.sol

@@ -4,8 +4,9 @@ pragma solidity ^0.8.13;
 import {Test, console} from "forge-std/Test.sol";
 import {MockERC20} from "solmate/test/utils/mocks/MockERC20.sol";
 import {ERC20} from "solmate/tokens/ERC20.sol";
+import {IAccessControl} from "@openzeppelin/contracts/access/IAccessControl.sol";
 
-import {RumpelWalletFactory} from "../src/RumpelWallet.sol";
+import {RumpelWalletFactory} from "../src/RumpelWalletFactory.sol";
 import {RumpelGuard} from "../src/RumpelGuard.sol";
 import {InitializationScript} from "../src/InitializationScript.sol";
 import {RumpelModule} from "../src/RumpelModule.sol";
@@ -134,6 +135,7 @@ contract RumpelWalletTest is Test {
             operation: Enum.Operation.Call
         });
 
+        // Sign tx for execution
         bytes32 txHash = safe.getTransactionHash(
             safeTX.to, safeTX.value, safeTX.data, safeTX.operation, 0, 0, 0, address(0), payable(address(0)), 0
         );
@@ -157,6 +159,37 @@ contract RumpelWalletTest is Test {
         assertEq(counter.count(), 1);
     }
 
+    function test_moduleAuth(address lad) public {
+        vm.assume(lad != admin);
+
+        address[] memory owners = new address[](1);
+        owners[0] = address(makeAddr("random 111")); // random owner
+
+        ISafe safe = ISafe(rumpelWalletFactory.createWallet(owners, 1, address(0), address(0), 0, payable(address(0))));
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IAccessControl.AccessControlUnauthorizedAccount.selector, lad, rumpelModule.DEFAULT_ADMIN_ROLE()
+            )
+        );
+        vm.prank(lad);
+        rumpelModule.exec(
+            safe, address(mockToken), 0, abi.encodeCall(ERC20.transfer, (RUMPEL_VAULT, 1.1e18)), Enum.Operation.Call
+        );
+    }
+
+    function test_guardAuth(address lad) public {
+        vm.assume(lad != admin);
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IAccessControl.AccessControlUnauthorizedAccount.selector, lad, rumpelModule.DEFAULT_ADMIN_ROLE()
+            )
+        );
+        vm.prank(lad);
+        rumpelGuard.setCallAllowed(address(counter), Counter.increment.selector, true);
+    }
+
     // test owner is blocked for actions not on the whitelist
     // - cant disable module
     // - cant change guard