From 0915fa33f1a52078aec4c60a3b5adc1bf5403b36 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Thu, 6 Feb 2025 14:34:48 +0100
Subject: [PATCH] vault: do not expose Lock internals on public api

---
 contracts/Vault.sol           |  9 +++++++--
 contracts/vault/VaultBase.sol | 13 ++++++++++---
 test/Vault.tests.js           | 23 ++++++++++++-----------
 test/vault.js                 |  8 ++++++++
 4 files changed, 37 insertions(+), 16 deletions(-)
 create mode 100644 test/vault.js

diff --git a/contracts/Vault.sol b/contracts/Vault.sol
index a9fc7d2..18533a8 100644
--- a/contracts/Vault.sol
+++ b/contracts/Vault.sol
@@ -25,9 +25,14 @@ contract Vault is VaultBase {
     return balance.designated;
   }
 
-  function getLock(Fund fund) public view returns (Lock memory) {
+  function getLockStatus(Fund fund) public view returns (LockStatus) {
     Controller controller = Controller.wrap(msg.sender);
-    return _getLock(controller, fund);
+    return _getLockStatus(controller, fund);
+  }
+
+  function getLockExpiry(Fund fund) public view returns (Timestamp) {
+    Controller controller = Controller.wrap(msg.sender);
+    return _getLockExpiry(controller, fund);
   }
 
   function lock(Fund fund, Timestamp expiry, Timestamp maximum) public {
diff --git a/contracts/vault/VaultBase.sol b/contracts/vault/VaultBase.sol
index 5207127..cb334ca 100644
--- a/contracts/vault/VaultBase.sol
+++ b/contracts/vault/VaultBase.sol
@@ -28,11 +28,18 @@ abstract contract VaultBase {
     _token = token;
   }
 
-  function _getLock(
+  function _getLockStatus(
     Controller controller,
     Fund fund
-  ) internal view returns (Lock memory) {
-    return _locks[controller][fund];
+  ) internal view returns (LockStatus) {
+    return _locks[controller][fund].status();
+  }
+
+  function _getLockExpiry(
+    Controller controller,
+    Fund fund
+  ) internal view returns (Timestamp) {
+    return _locks[controller][fund].expiry;
   }
 
   function _getBalance(
diff --git a/test/Vault.tests.js b/test/Vault.tests.js
index 13dcdcb..b153f65 100644
--- a/test/Vault.tests.js
+++ b/test/Vault.tests.js
@@ -10,6 +10,7 @@ const {
   snapshot,
   revert,
 } = require("./evm")
+const { LockStatus } = require("./vault")
 
 describe("Vault", function () {
   const fund = randomBytes(32)
@@ -46,8 +47,8 @@ describe("Vault", function () {
       expiry = (await currentTime()) + 80
       maximum = (await currentTime()) + 100
       await vault.lock(fund, expiry, maximum)
-      expect((await vault.getLock(fund))[0]).to.equal(expiry)
-      expect((await vault.getLock(fund))[1]).to.equal(maximum)
+      expect(await vault.getLockStatus(fund)).to.equal(LockStatus.Locked)
+      expect(await vault.getLockExpiry(fund)).to.equal(expiry)
     })
 
     it("does not allow a lock with expiry past maximum", async function () {
@@ -87,9 +88,9 @@ describe("Vault", function () {
 
       it("can extend a lock expiry up to its maximum", async function () {
         await vault.extendLock(fund, expiry + 1)
-        expect((await vault.getLock(fund))[0]).to.equal(expiry + 1)
+        expect(await vault.getLockExpiry(fund)).to.equal(expiry + 1)
         await vault.extendLock(fund, maximum)
-        expect((await vault.getLock(fund))[0]).to.equal(maximum)
+        expect(await vault.getLockExpiry(fund)).to.equal(maximum)
       })
 
       it("cannot extend a lock past its maximum", async function () {
@@ -106,8 +107,8 @@ describe("Vault", function () {
         await token.connect(controller).approve(vault.address, 30)
         await vault.deposit(fund, account.address, 30)
         await vault.burn(fund, account.address)
-        expect((await vault.getLock(fund))[0]).to.not.equal(0)
-        expect((await vault.getLock(fund))[1]).to.not.equal(0)
+        expect(await vault.getLockStatus(fund)).to.equal(LockStatus.Locked)
+        expect(await vault.getLockExpiry(fund)).to.not.equal(0)
       })
     })
 
@@ -587,15 +588,15 @@ describe("Vault", function () {
         await expire()
         // some tokens are withdrawn
         await vault.withdraw(fund, account.address)
-        expect((await vault.getLock(fund))[0]).not.to.equal(0)
-        expect((await vault.getLock(fund))[1]).not.to.equal(0)
+        expect(await vault.getLockStatus(fund)).to.equal(LockStatus.Unlocked)
+        expect(await vault.getLockExpiry(fund)).not.to.equal(0)
         // remainder of the tokens are withdrawn by recipient
         await vault
           .connect(account3)
           .withdrawByRecipient(controller.address, fund)
-        expect((await vault.getLock(fund))[0]).to.equal(0)
-        expect((await vault.getLock(fund))[1]).to.equal(0)
-      })
+        expect(await vault.getLockStatus(fund)).to.equal(LockStatus.NoLock)
+        expect(await vault.getLockExpiry(fund)).to.equal(0)
+        })
     })
 
     describe("flowing", function () {
diff --git a/test/vault.js b/test/vault.js
new file mode 100644
index 0000000..3d4f1a2
--- /dev/null
+++ b/test/vault.js
@@ -0,0 +1,8 @@
+const LockStatus = {
+  NoLock: 0,
+  Locked: 1,
+  Unlocked: 2,
+  Burned: 3
+}
+
+module.exports = { LockStatus }