feat: lock naked aura

src/AuraLockerModule.sol

@@ -5,7 +5,7 @@ import {IERC20} from "@openzeppelin/token/ERC20/IERC20.sol";
 
 import {KeeperCompatibleInterface} from "@chainlink/automation/interfaces/KeeperCompatibleInterface.sol";
 
-import {IGnosisSafe} from "./interfaces/gnosis/IGnosisSafe.sol";
+import {ISafe} from "./interfaces/gnosis/ISafe.sol";
 import {ILockAura} from "./interfaces/aura/ILockAura.sol";
 
 /// @title AuraLockerModule
@@ -17,8 +17,8 @@ contract AuraLockerModule is
     /*//////////////////////////////////////////////////////////////////////////
                                    CONSTANTS
     //////////////////////////////////////////////////////////////////////////*/
-    address public constant BALANCER_MULTISIG = 0x10A19e7eE7d7F8a52822f6817de8ea18204F2e4f;
-    IGnosisSafe public constant SAFE = IGnosisSafe(payable(BALANCER_MULTISIG));
+    address public constant BALANCER_GOV_SAFE = 0x9a5BDF08a6969A4bDb7724beE3c6d8964BDc0B28;
+    ISafe public constant SAFE = ISafe(payable(BALANCER_GOV_SAFE));
 
     IERC20 public constant AURA = IERC20(0xC0c293ce456fF0ED870ADd98a0828Dd4d2903DBF);
 
@@ -66,7 +66,7 @@ contract AuraLockerModule is
 
     /// @notice Enforce that the function is called by governance only
     modifier onlyGovernance() {
-        if (msg.sender != BALANCER_MULTISIG) revert NotGovernance(msg.sender);
+        if (msg.sender != BALANCER_GOV_SAFE) revert NotGovernance(msg.sender);
         _;
     }
 
@@ -94,38 +94,67 @@ contract AuraLockerModule is
         override
         returns (bool requiresLocking, bytes memory execPayload)
     {
-        if (!_isModuleEnabled()) return (false, bytes("AuraLocker module is not enabled"));
+        if (!SAFE.isModuleEnabled(address(this))) return (false, bytes("AuraLocker module is not enabled"));
 
-        (, uint256 unlockable,,) = AURA_LOCKER.lockedBalances(address(SAFE));
-
-        if (unlockable > 0) {
+        (, uint256 relockable,,) = AURA_LOCKER.lockedBalances(address(SAFE));
+        if (relockable > 0) {
             return (true, abi.encodeWithSelector(AURA_LOCKER.processExpiredLocks.selector, true));
         }
 
+        uint256 auraBalance = AURA.balanceOf(address(SAFE));
+        if (auraBalance > 0) {
+            return (true, abi.encodeWithSelector(AURA_LOCKER.lock.selector, address(SAFE), auraBalance));
+        }
+
         return (false, bytes("No AURA tokens unlocked"));
     }
 
     /// @notice The actual execution of the action determined by the `checkUpkeep` method (AURA locking)
     function performUpkeep(bytes calldata /* _performData */ ) external override onlyKeeper {
-        if (!_isModuleEnabled()) revert ModuleNotEnabled();
+        // Check if the module is enabled
+        if (SAFE.isModuleEnabled(address(this)) == false) {
+            revert ModuleNotEnabled();
+        }
 
-        (, uint256 unlockable,,) = AURA_LOCKER.lockedBalances(address(SAFE));
-        if (unlockable == 0) revert NothingToLock(block.timestamp);
+        // Relock expired locks if there are any
+        (, uint256 relockable,,) = AURA_LOCKER.lockedBalances(address(SAFE));
+        if (relockable > 0) {
+            // execute: `processExpiredLocks` via module
+            bool processExpiredLocksSucceeded = SAFE.execTransactionFromModule(
+                address(AURA_LOCKER), 0, abi.encodeCall(ILockAura.processExpiredLocks, true), ISafe.Operation.Call
+            );
+            if (processExpiredLocksSucceeded == false) {
+                revert TxFromModuleFailed();
+            }
+        }
 
-        // execute: `processExpiredLocks` via module
-        if (
-            !SAFE.execTransactionFromModule(
-                address(AURA_LOCKER), 0, abi.encodeCall(ILockAura.processExpiredLocks, true), IGnosisSafe.Operation.Call
-            )
-        ) revert TxFromModuleFailed();
-    }
+        // Lock AURA tokens if there are any
+        uint256 auraBalance = AURA.balanceOf(address(SAFE));
+        if (auraBalance > 0) {
+            // execute: `approve` via module
+            bool approveCallSucceeded = SAFE.execTransactionFromModule(
+                address(AURA),
+                0,
+                abi.encodeCall(IERC20.approve, (address(AURA_LOCKER), auraBalance)),
+                ISafe.Operation.Call
+            );
+            if (approveCallSucceeded == false) {
+                revert TxFromModuleFailed();
+            }
+            // execute: `lock` via module
+            bool lockCallSucceeded = SAFE.execTransactionFromModule(
+                address(AURA_LOCKER),
+                0,
+                abi.encodeCall(ILockAura.lock, (address(SAFE), auraBalance)),
+                ISafe.Operation.Call
+            );
+            if (lockCallSucceeded == false) {
+                revert TxFromModuleFailed();
+            }
+        }
 
-    /// @dev The Gnosis Safe v1.1.1 does not yet have the `isModuleEnabled` method, so we need a workaround
-    function _isModuleEnabled() internal view returns (bool) {
-        address[] memory modules = SAFE.getModules();
-        for (uint256 i = 0; i < modules.length; i++) {
-            if (modules[i] == address(this)) return true;
+        if (relockable == 0 && auraBalance == 0) {
+            revert NothingToLock(block.timestamp);
         }
-        return false;
     }
 }

src/interfaces/gnosis/IGnosisSafe.sol → src/interfaces/gnosis/ISafe.sol

@@ -1,34 +1,43 @@
-// SPDX-License-Identifier: MIT
+// SPDX-License-Identifier: LGPL-3.0-only
 pragma solidity ^0.8.4;
 
-interface IGnosisSafe {
+interface ISafe {
     enum Operation {
         Call,
         DelegateCall
     }
 
-    event AddedOwner(address owner);
+    event AddedOwner(address indexed owner);
     event ApproveHash(bytes32 indexed approvedHash, address indexed owner);
-    event ChangedMasterCopy(address masterCopy);
+    event ChangedFallbackHandler(address indexed handler);
+    event ChangedGuard(address indexed guard);
     event ChangedThreshold(uint256 threshold);
-    event DisabledModule(address module);
-    event EnabledModule(address module);
-    event ExecutionFailure(bytes32 txHash, uint256 payment);
+    event DisabledModule(address indexed module);
+    event EnabledModule(address indexed module);
+    event ExecutionFailure(bytes32 indexed txHash, uint256 payment);
     event ExecutionFromModuleFailure(address indexed module);
     event ExecutionFromModuleSuccess(address indexed module);
-    event ExecutionSuccess(bytes32 txHash, uint256 payment);
-    event RemovedOwner(address owner);
+    event ExecutionSuccess(bytes32 indexed txHash, uint256 payment);
+    event RemovedOwner(address indexed owner);
+    event SafeReceived(address indexed sender, uint256 value);
+    event SafeSetup(
+        address indexed initiator, address[] owners, uint256 threshold, address initializer, address fallbackHandler
+    );
     event SignMsg(bytes32 indexed msgHash);
 
-    fallback() external payable;
+    fallback() external;
+
+    receive() external payable;
 
-    function NAME() external view returns (string memory);
     function VERSION() external view returns (string memory);
     function addOwnerWithThreshold(address owner, uint256 _threshold) external;
     function approveHash(bytes32 hashToApprove) external;
     function approvedHashes(address, bytes32) external view returns (uint256);
-    function changeMasterCopy(address _masterCopy) external;
     function changeThreshold(uint256 _threshold) external;
+    function checkNSignatures(bytes32 dataHash, bytes memory data, bytes memory signatures, uint256 requiredSignatures)
+        external
+        view;
+    function checkSignatures(bytes32 dataHash, bytes memory data, bytes memory signatures) external view;
     function disableModule(address prevModule, address module) external;
     function domainSeparator() external view returns (bytes32);
     function enableModule(address module) external;
@@ -55,20 +64,20 @@ interface IGnosisSafe {
         address gasToken,
         address payable refundReceiver,
         bytes memory signatures
-    ) external returns (bool success);
+    ) external payable returns (bool success);
     function execTransactionFromModule(address to, uint256 value, bytes memory data, Operation operation)
         external
         returns (bool success);
     function execTransactionFromModuleReturnData(address to, uint256 value, bytes memory data, Operation operation)
         external
         returns (bool success, bytes memory returnData);
-    function getMessageHash(bytes memory message) external view returns (bytes32);
-    function getModules() external view returns (address[] memory);
+    function getChainId() external view returns (uint256);
     function getModulesPaginated(address start, uint256 pageSize)
         external
         view
         returns (address[] memory array, address next);
     function getOwners() external view returns (address[] memory);
+    function getStorageAt(uint256 offset, uint256 length) external view returns (bytes memory);
     function getThreshold() external view returns (uint256);
     function getTransactionHash(
         address to,
@@ -82,14 +91,12 @@ interface IGnosisSafe {
         address refundReceiver,
         uint256 _nonce
     ) external view returns (bytes32);
+    function isModuleEnabled(address module) external view returns (bool);
     function isOwner(address owner) external view returns (bool);
-    function isValidSignature(bytes memory _data, bytes memory _signature) external returns (bytes4);
     function nonce() external view returns (uint256);
     function removeOwner(address prevOwner, address owner, uint256 _threshold) external;
-    function requiredTxGas(address to, uint256 value, bytes memory data, Operation operation)
-        external
-        returns (uint256);
     function setFallbackHandler(address handler) external;
+    function setGuard(address guard) external;
     function setup(
         address[] memory _owners,
         uint256 _threshold,
@@ -100,7 +107,7 @@ interface IGnosisSafe {
         uint256 payment,
         address payable paymentReceiver
     ) external;
-    function signMessage(bytes memory _data) external;
     function signedMessages(bytes32) external view returns (uint256);
+    function simulateAndRevert(address targetContract, bytes memory calldataPayload) external;
     function swapOwner(address prevOwner, address oldOwner, address newOwner) external;
 }

test/AuraLockerModule.t.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.25;
 
 import {BaseFixture} from "./BaseFixture.sol";
 
+import {IERC20} from "@openzeppelin/token/ERC20/IERC20.sol";
 import {ILockAura} from "../src/interfaces/aura/ILockAura.sol";
 
 import {AuraLockerModule} from "../src/AuraLockerModule.sol";
@@ -27,10 +28,7 @@ contract AuraLockerModuleTest is BaseFixture {
         // `disableModule(address prevModule, address module)`
         vm.prank(address(SAFE));
         SAFE.disableModule(address(1), address(auraLockerModule));
-        address[] memory modules = SAFE.getModules();
-        for (uint256 i = 0; i < modules.length; i++) {
-            if (modules[i] == address(auraLockerModule)) assertFalse(true);
-        }
+        assertFalse(SAFE.isModuleEnabled(address(auraLockerModule)));
 
         // once module is removed, the keeper trying to call `performUpkeep` should revert
         vm.prank(auraLockerModule.keeper());
@@ -81,4 +79,46 @@ contract AuraLockerModuleTest is BaseFixture {
         vm.expectRevert(abi.encodeWithSelector(AuraLockerModule.ZeroAddressValue.selector));
         auraLockerModule.setKeeper(address(0));
     }
+
+    function testAutomaticLockingOfNakedAura() public {
+        // Get AURA token reference
+        IERC20 aura = IERC20(0xC0c293ce456fF0ED870ADd98a0828Dd4d2903DBF);
+
+        // Use a known AURA holder address from mainnet
+        // This is a treasury or large holder address with sufficient AURA
+        address auraWhale = 0x43B17088503F4CE1AED9fB302ED6BB51aD6694Fa; // Aura Treasury
+
+        uint256 whaleBalance = aura.balanceOf(auraWhale);
+
+        // Ensure we have enough balance to transfer
+        assertGe(whaleBalance, 1e18, "Whale doesn't have enough AURA");
+
+        // Transfer 1 AURA to the safe
+        vm.prank(auraWhale);
+        aura.transfer(address(SAFE), 1e18);
+
+        // Verify AURA was received
+        uint256 safeBalance = aura.balanceOf(address(SAFE));
+        assertEq(safeBalance, 1e18, "Safe should have 1 AURA");
+
+        // Check that upkeep is needed
+        (bool requiresLocking, bytes memory execPayload) = auraLockerModule.checkUpkeep(bytes(""));
+        assertTrue(requiresLocking, "Should require locking");
+        assertEq(execPayload, abi.encodeWithSelector(AURA_LOCKER.lock.selector, address(SAFE), 1e18));
+
+        // Get current locked balance
+        (uint256 totalLockedBefore,, uint256 lockedBefore,) = AURA_LOCKER.lockedBalances(address(SAFE));
+
+        // Perform the upkeep
+        vm.prank(auraLockerModule.keeper());
+        auraLockerModule.performUpkeep(bytes(""));
+
+        // Verify AURA was locked
+        (uint256 totalLockedAfter,, uint256 lockedAfter,) = AURA_LOCKER.lockedBalances(address(SAFE));
+        assertEq(totalLockedAfter, totalLockedBefore + 1e18, "Total locked should increase by 1 AURA");
+        assertEq(lockedAfter, lockedBefore + 1e18, "Locked balance should increase by 1 AURA");
+
+        // Verify safe no longer has AURA
+        assertEq(aura.balanceOf(address(SAFE)), 0, "Safe should have 0 AURA after locking");
+    }
 }

test/BaseFixture.sol

@@ -8,7 +8,7 @@ import {IERC20} from "@openzeppelin/token/ERC20/IERC20.sol";
 import {IKeeperRegistryMaster} from "@chainlink/automation/interfaces/v2_1/IKeeperRegistryMaster.sol";
 import {IKeeperRegistrar} from "../src/interfaces/chainlink/IKeeperRegistrar.sol";
 
-import {IGnosisSafe} from "../src/interfaces/gnosis/IGnosisSafe.sol";
+import {ISafe} from "../src/interfaces/gnosis/ISafe.sol";
 import {ILockAura} from "../src/interfaces/aura/ILockAura.sol";
 
 import {AuraLockerModule} from "../src/AuraLockerModule.sol";
@@ -18,7 +18,7 @@ contract BaseFixture is Test {
     // https://debank.com/profile/0x9ff471F9f98F42E5151C7855fD1b5aa906b1AF7e
     address constant BALANCER_ADMIN_CHAINLINK_UPKEEPS = 0x9ff471F9f98F42E5151C7855fD1b5aa906b1AF7e;
 
-    IGnosisSafe public constant SAFE = IGnosisSafe(payable(0x10A19e7eE7d7F8a52822f6817de8ea18204F2e4f));
+    ISafe public constant SAFE = ISafe(payable(0x9a5BDF08a6969A4bDb7724beE3c6d8964BDc0B28));
 
     // https://docs.chain.link/resources/link-token-contracts?parent=automation#ethereum-mainnet
     IERC20 constant LINK = IERC20(0x514910771AF9Ca656af840dff83E8264EcF986CA);
@@ -31,8 +31,8 @@ contract BaseFixture is Test {
     AuraLockerModule auraLockerModule;
 
     function setUp() public virtual {
-        // block @ https://etherscan.io/block/20576471
-        vm.createSelectFork("ethereum", 20576471);
+        // block @ https://etherscan.io/block/22686870
+        vm.createSelectFork("ethereum", 22686870);
 
         // deploy module
         auraLockerModule = new AuraLockerModule();
@@ -72,7 +72,7 @@ contract BaseFixture is Test {
         vm.label(address(AURA_LOCKER), "AURA_LOCKER");
         vm.label(address(auraLockerModule), "AURA_LOCKER_MODULE");
         vm.label(BALANCER_ADMIN_CHAINLINK_UPKEEPS, "BALANCER_ADMIN_CHAINLINK_UPKEEPS");
-        vm.label(address(SAFE), "BALANCER_MULTISIG");
+        vm.label(address(SAFE), "BALANCER_GOV_SAFE");
         vm.label(address(CL_REGISTRY), "CL_REGISTRY");
         vm.label(address(CL_REGISTRAR), "CL_REGISTRAR");
     }