From d35cc0084151e3c410c76906da78061d77a679dd Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Tue, 28 Jan 2025 15:00:20 +0100
Subject: [PATCH] vault: check Lock invariant before writing

---
 contracts/vault/VaultBase.sol | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/contracts/vault/VaultBase.sol b/contracts/vault/VaultBase.sol
index 6513dc5..9016f46 100644
--- a/contracts/vault/VaultBase.sol
+++ b/contracts/vault/VaultBase.sol
@@ -149,10 +149,12 @@ abstract contract VaultBase {
     Timestamp expiry,
     Timestamp maximum
   ) internal {
-    require(expiry <= maximum, ExpiryPastMaximum());
-    Lock memory existing = _locks[controller][context];
-    require(existing.maximum == Timestamp.wrap(0), AlreadyLocked());
-    _locks[controller][context] = Lock({expiry: expiry, maximum: maximum});
+    Lock memory lock = _locks[controller][context];
+    require(lock.maximum == Timestamp.wrap(0), AlreadyLocked());
+    lock.expiry = expiry;
+    lock.maximum = maximum;
+    _checkLockInvariant(lock);
+    _locks[controller][context] = lock;
   }
 
   function _extendLock(
@@ -163,8 +165,9 @@ abstract contract VaultBase {
     Lock memory lock = _locks[controller][context];
     require(lock.isLocked(), LockRequired());
     require(lock.expiry <= expiry, InvalidExpiry());
-    require(expiry <= lock.maximum, ExpiryPastMaximum());
-    _locks[controller][context].expiry = expiry;
+    lock.expiry = expiry;
+    _checkLockInvariant(lock);
+    _locks[controller][context] = lock;
   }
 
   function _flow(
@@ -198,6 +201,10 @@ abstract contract VaultBase {
     _flows[controller][context][to] = receiverFlow;
   }
 
+  function _checkLockInvariant(Lock memory lock) private pure {
+    require(lock.expiry <= lock.maximum, ExpiryPastMaximum());
+  }
+
   function _checkFlowInvariant(
     Balance memory balance,
     Lock memory lock,