From 91a976a00722f06e5d1f5be58d242547060c8646 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Tue, 15 Feb 2022 17:18:25 +0100
Subject: [PATCH] Prevent withdrawal of locked collateral

---
 contracts/Collateral.sol     |  4 +++-
 contracts/TestCollateral.sol | 12 ++++++++++++
 test/Collateral.test.js      | 22 ++++++++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/contracts/Collateral.sol b/contracts/Collateral.sol
index e3087ce..f65df55 100644
--- a/contracts/Collateral.sol
+++ b/contracts/Collateral.sol
@@ -2,8 +2,9 @@
 pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "./AccountLocks.sol";
 
-contract Collateral {
+contract Collateral is AccountLocks {
   IERC20 private immutable token;
   Totals private totals;
   mapping(address => uint256) private balances;
@@ -38,6 +39,7 @@ contract Collateral {
   }
 
   function withdraw() public invariant {
+    _unlockAccount();
     uint256 amount = balanceOf(msg.sender);
     totals.withdrawn += amount;
     subtract(msg.sender, amount);
diff --git a/contracts/TestCollateral.sol b/contracts/TestCollateral.sol
index 5f5cb4a..2524c76 100644
--- a/contracts/TestCollateral.sol
+++ b/contracts/TestCollateral.sol
@@ -11,4 +11,16 @@ contract TestCollateral is Collateral {
   function slash(address account, uint256 percentage) public {
     _slash(account, percentage);
   }
+
+  function createLock(bytes32 id, uint256 expiry) public {
+    _createLock(id, expiry);
+  }
+
+  function lock(address account, bytes32 id) public {
+    _lock(account, id);
+  }
+
+  function unlock(bytes32 id) public {
+    _unlock(id);
+  }
 }
diff --git a/test/Collateral.test.js b/test/Collateral.test.js
index c747a73..0fdab4f 100644
--- a/test/Collateral.test.js
+++ b/test/Collateral.test.js
@@ -1,4 +1,5 @@
 const { expect } = require("chai")
+const { exampleLock } = require("./examples")
 
 describe("Collateral", function () {
   let collateral, token
@@ -89,4 +90,25 @@ describe("Collateral", function () {
       expect(await collateral.balanceOf(account1.address)).to.equal(950)
     })
   })
+
+  describe("locking", function () {
+    let lock
+
+    beforeEach(async function () {
+      await token.approve(collateral.address, 42)
+      await collateral.deposit(42)
+      lock = exampleLock()
+      await collateral.createLock(lock.id, lock.expiry)
+      await collateral.lock(account0.address, lock.id)
+    })
+
+    it("withdrawal fails when account is locked", async function () {
+      await expect(collateral.withdraw()).to.be.revertedWith("Account locked")
+    })
+
+    it("withdrawal succeeds when account is unlocked", async function () {
+      await collateral.unlock(lock.id)
+      await expect(collateral.withdraw()).not.to.be.reverted
+    })
+  })
 })