Fix vault lock calculation timing to ensure consistency with StakeManager

Co-authored-by: 3esmit <224810+3esmit@users.noreply.github.com>

src/StakeVault.sol

@@ -151,13 +151,16 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
      * @param _seconds The time period to lock the staked amount for.
      */
     function lock(uint256 _seconds) external onlyOwner onlyTrustedStakeManager {
-        // Update lock time before calling stake manager, using same logic as StakeManager
-        if (_seconds > 0) {
-            uint256 newLockEnd = Math.max(lockUntil, block.timestamp) + _seconds;
-            lockUntil = newLockEnd;
-        }
+        // Store old lock time for calculation
+        uint256 oldLockUntil = lockUntil;
         
         stakeManager.lock(_seconds);
+        
+        // Update lock time after manager call, using same logic as StakeManager
+        if (_seconds > 0) {
+            uint256 newLockEnd = Math.max(oldLockUntil, block.timestamp) + _seconds;
+            lockUntil = newLockEnd;
+        }
     }
 
     /**
@@ -327,13 +330,17 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
      * @param _source The address from which tokens will be transferred.
      */
     function _stake(uint256 _amount, uint256 _seconds, address _source) internal {
-        // Update lock time before calling stake manager, using same logic as StakeManager
+        // Store old lock time for calculation
+        uint256 oldLockUntil = lockUntil;
+        
+        stakeManager.stake(_amount, _seconds);
+        
+        // Update lock time after manager call, using same logic as StakeManager
         if (_seconds > 0) {
-            uint256 newLockEnd = Math.max(lockUntil, block.timestamp) + _seconds;
+            uint256 newLockEnd = Math.max(oldLockUntil, block.timestamp) + _seconds;
             lockUntil = newLockEnd;
         }
         
-        stakeManager.stake(_amount, _seconds);
         bool success = STAKING_TOKEN.transferFrom(_source, address(this), _amount);
         if (!success) {
             revert StakeVault__StakingFailed();

src/interfaces/IStakeVault.sol

@@ -8,5 +8,6 @@ interface IStakeVault {
     function stakeManager() external view returns (IStakeManagerProxy);
     function register() external;
     function lockUntil() external view returns (uint256);
+    /// @notice Updates lock time - primarily used for vault migration
     function updateLockUntil(uint256 newLockUntil) external;
 }

test/LockLogic.t.sol

@@ -0,0 +1,111 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.26;
+
+import { Test } from "forge-std/Test.sol";
+import { StakeVault } from "../src/StakeVault.sol";
+import { MockStakeManager } from "./mocks/MockStakeManager.sol";
+import { MockToken } from "./mocks/MockToken.sol";
+
+/**
+ * @title LockLogicTest
+ * @notice Tests to verify the new vault-driven lock logic works correctly
+ */
+contract LockLogicTest is Test {
+    StakeVault internal stakeVault;
+    MockStakeManager internal stakeManager;
+    MockToken internal stakingToken;
+    address internal alice = makeAddr("alice");
+
+    function setUp() public {
+        stakingToken = new MockToken("Staking Token", "ST");
+        stakeManager = new MockStakeManager();
+        
+        // Create vault directly (without factory for simplicity)
+        stakeVault = new StakeVault(stakingToken);
+        stakeVault.initialize(alice, address(stakeManager));
+        
+        // Mint tokens to alice and approve vault
+        stakingToken.mint(alice, 10_000e18);
+        vm.prank(alice);
+        stakingToken.approve(address(stakeVault), 10_000e18);
+    }
+
+    function test_StakeUpdatesLockUntilCorrectly() public {
+        // Initial lockUntil should be 0
+        assertEq(stakeVault.lockUntil(), 0);
+        
+        // Stake with 90 days lock
+        uint256 lockPeriod = 90 days;
+        uint256 stakeTime = block.timestamp;
+        
+        vm.prank(alice);
+        stakeVault.stake(1000e18, lockPeriod);
+        
+        // Vault should have updated its lockUntil before calling manager
+        uint256 expectedLockEnd = stakeTime + lockPeriod; // since initial lockUntil was 0
+        assertEq(stakeVault.lockUntil(), expectedLockEnd);
+    }
+
+    function test_StakeExtendsExistingLock() public {
+        // Set initial lock time in the future
+        uint256 initialLockEnd = block.timestamp + 60 days;
+        vm.prank(address(stakeManager));
+        stakeVault.updateLockUntil(initialLockEnd);
+        
+        // Stake with additional 90 days
+        uint256 additionalLockPeriod = 90 days;
+        uint256 stakeTime = block.timestamp;
+        
+        vm.prank(alice);
+        stakeVault.stake(1000e18, additionalLockPeriod);
+        
+        // Should extend from existing lock, not from current time
+        uint256 expectedLockEnd = initialLockEnd + additionalLockPeriod;
+        assertEq(stakeVault.lockUntil(), expectedLockEnd);
+    }
+
+    function test_StakeExtendsFromCurrentTimeIfLockExpired() public {
+        // Set initial lock time in the past
+        uint256 pastLockEnd = block.timestamp - 30 days;
+        vm.prank(address(stakeManager));
+        stakeVault.updateLockUntil(pastLockEnd);
+        
+        // Stake with 90 days lock
+        uint256 lockPeriod = 90 days;
+        uint256 stakeTime = block.timestamp;
+        
+        vm.prank(alice);
+        stakeVault.stake(1000e18, lockPeriod);
+        
+        // Should extend from current time since old lock expired
+        uint256 expectedLockEnd = stakeTime + lockPeriod;
+        assertEq(stakeVault.lockUntil(), expectedLockEnd);
+    }
+
+    function test_LockUpdatesLockUntilCorrectly() public {
+        // Initial stake without lock
+        vm.prank(alice);
+        stakeVault.stake(1000e18, 0);
+        
+        assertEq(stakeVault.lockUntil(), 0);
+        
+        // Now lock for 90 days
+        uint256 lockPeriod = 90 days;
+        uint256 lockTime = block.timestamp;
+        
+        vm.prank(alice);
+        stakeVault.lock(lockPeriod);
+        
+        uint256 expectedLockEnd = lockTime + lockPeriod;
+        assertEq(stakeVault.lockUntil(), expectedLockEnd);
+    }
+
+    function test_StakeWithZeroLockDoesNotUpdateLockUntil() public {
+        // Stake without lock period
+        vm.prank(alice);
+        stakeVault.stake(1000e18, 0);
+        
+        // lockUntil should remain 0  
+        assertEq(stakeVault.lockUntil(), 0);
+    }
+}

test/MigrationLogic.t.sol

@@ -0,0 +1,61 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.26;
+
+import { Test } from "forge-std/Test.sol";
+import { StakeVault } from "../src/StakeVault.sol";
+import { MockStakeManager } from "./mocks/MockStakeManager.sol";
+import { MockToken } from "./mocks/MockToken.sol";
+
+/**
+ * @title MigrationLogicTest  
+ * @notice Tests to verify migration still works with new lock logic
+ */
+contract MigrationLogicTest is Test {
+    StakeVault internal sourceVault;
+    StakeVault internal targetVault;
+    MockStakeManager internal stakeManager;
+    MockToken internal stakingToken;
+    address internal alice = makeAddr("alice");
+
+    function setUp() public {
+        stakingToken = new MockToken("Staking Token", "ST");
+        stakeManager = new MockStakeManager();
+        
+        // Create two vaults for migration test
+        sourceVault = new StakeVault(stakingToken);
+        sourceVault.initialize(alice, address(stakeManager));
+        
+        targetVault = new StakeVault(stakingToken);
+        targetVault.initialize(alice, address(stakeManager));
+        
+        // Mint tokens to alice and approve source vault
+        stakingToken.mint(alice, 10_000e18);
+        vm.prank(alice);
+        stakingToken.approve(address(sourceVault), 10_000e18);
+    }
+
+    function test_UpdateLockUntilStillWorksForMigration() public {
+        // Set a lock time on the source vault  
+        uint256 lockTime = block.timestamp + 90 days;
+        vm.prank(address(stakeManager));
+        sourceVault.updateLockUntil(lockTime);
+        
+        assertEq(sourceVault.lockUntil(), lockTime);
+        assertEq(targetVault.lockUntil(), 0);
+        
+        // Migration should be able to transfer lock time to target vault
+        vm.prank(address(stakeManager));
+        targetVault.updateLockUntil(lockTime);
+        
+        assertEq(targetVault.lockUntil(), lockTime);
+    }
+
+    function test_UpdateLockUntilOnlyWorksFromStakeManager() public {
+        uint256 lockTime = block.timestamp + 90 days;
+        
+        // Should revert when called by non-stake-manager
+        vm.prank(alice);
+        vm.expectRevert(StakeVault.StakeVault__StakeManagerImplementationNotTrusted.selector);
+        sourceVault.updateLockUntil(lockTime);
+    }
+}