vault: lock up tokens until expiry time

contracts/Timestamps.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.28;
+
+type Timestamp is uint64;
+
+library Timestamps {
+  function isAfter(Timestamp a, Timestamp b) internal pure returns (bool) {
+    return Timestamp.unwrap(a) > Timestamp.unwrap(b);
+  }
+
+  function isFuture(Timestamp timestamp) internal view returns (bool) {
+    return Timestamp.unwrap(timestamp) > block.timestamp;
+  }
+}

contracts/Vault.sol

@@ -3,8 +3,10 @@ pragma solidity 0.8.28;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "./Timestamps.sol";
 
 using SafeERC20 for IERC20;
+using Timestamps for Timestamp;
 
 contract Vault {
   IERC20 private immutable _token;
@@ -13,6 +15,12 @@ contract Vault {
   type Context is bytes32;
   type Recipient is address;
 
+  struct Lock {
+    Timestamp expiry;
+    Timestamp maximum;
+  }
+
+  mapping(Controller => mapping(Context => Lock)) private _locks;
   mapping(Controller => mapping(Context => mapping(Recipient => uint256)))
     private _available;
   mapping(Controller => mapping(Context => mapping(Recipient => uint256)))
@@ -40,6 +48,11 @@ contract Vault {
     return _designated[controller][context][recipient];
   }
 
+  function lock(Context context) public view returns (Lock memory) {
+    Controller controller = Controller.wrap(msg.sender);
+    return _locks[controller][context];
+  }
+
   function deposit(Context context, address from, uint256 amount) public {
     Controller controller = Controller.wrap(msg.sender);
     Recipient recipient = Recipient.wrap(from);
@@ -54,6 +67,7 @@ contract Vault {
   }
 
   function withdraw(Context context, Recipient recipient) public {
+    require(!lock(context).expiry.isFuture(), Locked());
     uint256 amount = balance(context, recipient);
     _delete(context, recipient);
     _token.safeTransfer(Recipient.unwrap(recipient), amount);
@@ -94,5 +108,26 @@ contract Vault {
     _designated[controller][context][recipient] += amount;
   }
 
+  function lockup(Context context, Timestamp expiry, Timestamp maximum) public {
+    require(Timestamp.unwrap(lock(context).maximum) == 0, AlreadyLocked());
+    require(!expiry.isAfter(maximum), ExpiryPastMaximum());
+    Controller controller = Controller.wrap(msg.sender);
+    _locks[controller][context] = Lock({expiry: expiry, maximum: maximum});
+  }
+
+  function extend(Context context, Timestamp expiry) public {
+    Lock memory previous = lock(context);
+    require(previous.expiry.isFuture(), LockExpired());
+    require(!previous.expiry.isAfter(expiry), InvalidExpiry());
+    require(!expiry.isAfter(previous.maximum), ExpiryPastMaximum());
+    Controller controller = Controller.wrap(msg.sender);
+    _locks[controller][context].expiry = expiry;
+  }
+
   error InsufficientBalance();
+  error Locked();
+  error AlreadyLocked();
+  error ExpiryPastMaximum();
+  error InvalidExpiry();
+  error LockExpired();
 }

test/Vault.tests.js

@@ -1,6 +1,7 @@
 const { expect } = require("chai")
 const { ethers } = require("hardhat")
 const { randomBytes } = ethers.utils
+const { currentTime, advanceTimeTo } = require("./evm")
 
 describe("Vault", function () {
   let token
@@ -264,4 +265,99 @@ describe("Vault", function () {
       expect(after - before).to.equal(amount)
     })
   })
+
+  describe("locking", async function () {
+    const context = randomBytes(32)
+    const amount = 42
+
+    beforeEach(async function () {
+      await token.connect(account).approve(vault.address, amount)
+      await vault.deposit(context, account.address, amount)
+    })
+
+    it("can lock up all tokens in a context", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      let maximum = start + 20
+      await vault.lockup(context, expiry, maximum)
+      expect((await vault.lock(context))[0]).to.equal(expiry)
+      expect((await vault.lock(context))[1]).to.equal(maximum)
+    })
+
+    it("cannot lock up when already locked", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      let maximum = start + 20
+      await vault.lockup(context, expiry, maximum)
+      const locking = vault.lockup(context, expiry, maximum)
+      await expect(locking).to.be.revertedWith("AlreadyLocked")
+    })
+
+    it("cannot lock when expiry is past maximum", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      let maximum = start + 9
+      const locking = vault.lockup(context, expiry, maximum)
+      await expect(locking).to.be.revertedWith("ExpiryPastMaximum")
+    })
+
+    it("does not allow withdrawal before lock expires", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      await vault.lockup(context, expiry, expiry)
+      await advanceTimeTo(expiry - 1)
+      const withdrawing = vault.withdraw(context, account.address)
+      await expect(withdrawing).to.be.revertedWith("Locked")
+    })
+
+    it("allows withdrawal after lock expires", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      await vault.lockup(context, expiry, expiry)
+      await advanceTimeTo(expiry)
+      const before = await token.balanceOf(account.address)
+      await vault.withdraw(context, account.address)
+      const after = await token.balanceOf(account.address)
+      expect(after - before).to.equal(amount)
+    })
+
+    it("can extend a lock expiry up to its maximum", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      let maximum = start + 20
+      await vault.lockup(context, expiry, maximum)
+      await vault.extend(context, start + 15)
+      expect((await vault.lock(context))[0]).to.equal(start + 15)
+      await vault.extend(context, start + 20)
+      expect((await vault.lock(context))[0]).to.equal(start + 20)
+    })
+
+    it("cannot extend a lock past its maximum", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      let maximum = start + 20
+      await vault.lockup(context, expiry, maximum)
+      const extending = vault.extend(context, start + 21)
+      await expect(extending).to.be.revertedWith("ExpiryPastMaximum")
+    })
+
+    it("cannot move expiry forward", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      let maximum = start + 20
+      await vault.lockup(context, expiry, maximum)
+      const extending = vault.extend(context, start + 9)
+      await expect(extending).to.be.revertedWith("InvalidExpiry")
+    })
+
+    it("cannot extend an expired lock", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      let maximum = start + 20
+      await vault.lockup(context, expiry, maximum)
+      await advanceTimeTo(expiry)
+      const extending = vault.extend(context, maximum)
+      await expect(extending).to.be.revertedWith("LockExpired")
+    })
+  })
 })