diff --git a/src/AuraLockerModule.sol b/src/AuraLockerModule.sol index 4f2ec86..41e5c8b 100644 --- a/src/AuraLockerModule.sol +++ b/src/AuraLockerModule.sol @@ -96,11 +96,22 @@ contract AuraLockerModule is { if (!SAFE.isModuleEnabled(address(this))) return (false, bytes("AuraLocker module is not enabled")); - (, uint256 relockable,,) = AURA_LOCKER.lockedBalances(address(SAFE)); + (, uint256 relockable,, ILockAura.LockedBalance[] memory lockData) = AURA_LOCKER.lockedBalances(address(SAFE)); if (relockable > 0) { return (true, abi.encodeWithSelector(AURA_LOCKER.processExpiredLocks.selector, true)); } + // Check if any locks are expiring within the next week + uint256 len = lockData.length; + if (len > 0) { + uint256 timestamp = block.timestamp; + for (uint256 i = 0; i < len; i++) { + if (timestamp + 1 weeks >= lockData[i].unlockTime) { + 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)); @@ -116,9 +127,22 @@ contract AuraLockerModule is revert ModuleNotEnabled(); } - // Relock expired locks if there are any - (, uint256 relockable,,) = AURA_LOCKER.lockedBalances(address(SAFE)); - if (relockable > 0) { + // Check if there are any expired locks + (, uint256 relockable,, ILockAura.LockedBalance[] memory lockData) = AURA_LOCKER.lockedBalances(address(SAFE)); + bool shouldRelock = relockable > 0; + + // Check if there are locks expiring soon + if (!shouldRelock && lockData.length > 0) { + uint256 timestamp = block.timestamp; + for (uint256 i = 0; i < lockData.length; i++) { + if (timestamp + 1 weeks >= lockData[i].unlockTime) { + shouldRelock = true; + break; + } + } + } + + if (shouldRelock) { // execute: `processExpiredLocks` via module bool processExpiredLocksSucceeded = SAFE.execTransactionFromModule( address(AURA_LOCKER), 0, abi.encodeCall(ILockAura.processExpiredLocks, true), ISafe.Operation.Call @@ -153,7 +177,7 @@ contract AuraLockerModule is } } - if (relockable == 0 && auraBalance == 0) { + if (!shouldRelock && auraBalance == 0) { revert NothingToLock(block.timestamp); } } diff --git a/test/AuraLockerModule.t.sol b/test/AuraLockerModule.t.sol index 5f5cc58..98c15e9 100644 --- a/test/AuraLockerModule.t.sol +++ b/test/AuraLockerModule.t.sol @@ -121,4 +121,57 @@ contract AuraLockerModuleTest is BaseFixture { // Verify safe no longer has AURA assertEq(aura.balanceOf(address(SAFE)), 0, "Safe should have 0 AURA after locking"); } + + function test_checkUpkeep_when_LocksExpiringSoon() public { + // skip forward to get closer to lock expiry + // existing lock expires around week 16 from the fork block + // we want to test the early trigger (within 1 week of expiry) + skip(10 weeks); + + // check current state + (, uint256 relockable,, ILockAura.LockedBalance[] memory lockData) = AURA_LOCKER.lockedBalances(address(SAFE)); + + // ensure we have locks that are not expired yet + assertEq(relockable, 0, "Should have no expired locks at this point"); + assertGt(lockData.length, 0, "Should have active locks"); + + // check that the module detects locks expiring within 1 week + (bool requiresLocking, bytes memory execPayload) = auraLockerModule.checkUpkeep(bytes("")); + + // verify the unlock time is within 1 week + uint256 timeUntilUnlock = lockData[0].unlockTime - block.timestamp; + if (timeUntilUnlock <= 1 weeks) { + assertTrue(requiresLocking, "Should require locking when locks expire within 1 week"); + assertEq(execPayload, abi.encodeWithSelector(AURA_LOCKER.processExpiredLocks.selector, true)); + } else { + assertFalse(requiresLocking, "Should not require locking when locks don't expire within 1 week"); + } + } + + function testPerformUpkeep_when_LocksExpiringSoon() public { + // skip to a point where locks will expire within 1 week but have not expired yet + skip(10 weeks); + + // check current state + (, uint256 relockable,, ILockAura.LockedBalance[] memory lockData) = AURA_LOCKER.lockedBalances(address(SAFE)); + assertEq(relockable, 0, "Should have no expired locks"); + assertGt(lockData.length, 0, "Should have active locks"); + + uint256 timeUntilUnlock = lockData[0].unlockTime - block.timestamp; + + // only test if locks are expiring within 1 week + if (timeUntilUnlock <= 1 weeks) { + // get the locked balance before + (uint256 totalBefore,, uint256 lockedBefore,) = AURA_LOCKER.lockedBalances(address(SAFE)); + + // perform upkeep + vm.prank(auraLockerModule.keeper()); + auraLockerModule.performUpkeep(bytes("")); + + // after performUpkeep, total and locked balances should still be the same + (uint256 totalAfter,, uint256 lockedAfter,) = AURA_LOCKER.lockedBalances(address(SAFE)); + assertEq(totalAfter, totalBefore, "Total AURA should remain the same"); + assertEq(lockedAfter, lockedBefore, "Locked amount should remain the same"); + } + } }