From 3d50e56fe7bb548c358f78abc9443fe27a2662f9 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Mon, 14 Feb 2022 16:19:47 +0100
Subject: [PATCH] Collateral invariant

---
 contracts/Collateral.sol | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/contracts/Collateral.sol b/contracts/Collateral.sol
index c09b907..cd29b65 100644
--- a/contracts/Collateral.sol
+++ b/contracts/Collateral.sol
@@ -7,7 +7,10 @@ contract Collateral {
   IERC20 private immutable token;
   mapping(address => uint256) private balances;
 
-  constructor(IERC20 _token) {
+  uint256 private totalDeposited;
+  uint256 private totalBalance;
+
+  constructor(IERC20 _token) invariant {
     token = _token;
   }
 
@@ -15,8 +18,15 @@ contract Collateral {
     return balances[account];
   }
 
-  function deposit(uint256 amount) public {
+  function deposit(uint256 amount) public invariant {
     token.transferFrom(msg.sender, address(this), amount);
+    totalDeposited += amount;
     balances[msg.sender] += amount;
+    totalBalance += amount;
+  }
+
+  modifier invariant() {
+    _;
+    assert(totalDeposited == totalBalance);
   }
 }