From 0d23116f370c362c7a85caf0671a81a5e8c4784e Mon Sep 17 00:00:00 2001
From: 0xyolo <yolo@cream.finance>
Date: Thu, 8 Apr 2021 23:26:47 +0800
Subject: [PATCH] update: totalBorrows

---
 contracts/CCapableErc20.sol           | 12 +++++++++---
 tests/Contracts/FlashloanReceiver.sol |  5 +++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/contracts/CCapableErc20.sol b/contracts/CCapableErc20.sol
index 0eeb54d02..f9e0552e3 100644
--- a/contracts/CCapableErc20.sol
+++ b/contracts/CCapableErc20.sol
@@ -229,6 +229,8 @@ contract CCapableErc20 is CToken, CCapableErc20Interface, CCapableDelegateInterf
     }
 
     function flashLoan(address receiver, uint amount, bytes calldata params) external nonReentrant {
+        require(accrueInterest() == uint(Error.NO_ERROR), "accrue interest failed");
+
         uint cashOnChainBefore = getCashOnChain();
         uint cashBefore = getCashPrior();
         require(cashBefore >= amount, "INSUFFICIENT_LIQUIDITY");
@@ -239,17 +241,21 @@ contract CCapableErc20 is CToken, CCapableErc20Interface, CCapableDelegateInterf
         // 2. transfer fund to receiver
         doTransferOut(address(uint160(receiver)), amount);
 
-        // 3. execute receiver's callback function
+        // 3. update totalBorrows
+        totalBorrows = add_(totalBorrows, amount);
+
+        // 4. execute receiver's callback function
         IFlashloanReceiver(receiver).executeOperation(msg.sender, underlying, amount, totalFee, params);
 
-        // 4. check balance
+        // 5. check balance
         uint cashOnChainAfter = getCashOnChain();
         require(cashOnChainAfter == add_(cashOnChainBefore, totalFee), "BALANCE_INCONSISTENT");
 
-        // 5. update reserves and internal cash
+        // 6. update reserves and internal cash and totalBorrows
         uint reservesFee = mul_ScalarTruncate(Exp({mantissa: reserveFactorMantissa}), totalFee);
         totalReserves = add_(totalReserves, reservesFee);
         internalCash = add_(cashBefore, totalFee);
+        totalBorrows = sub_(totalBorrows, amount);
 
         emit Flashloan(receiver, amount, totalFee, reservesFee);
     }
diff --git a/tests/Contracts/FlashloanReceiver.sol b/tests/Contracts/FlashloanReceiver.sol
index 2e0716183..60579dd58 100644
--- a/tests/Contracts/FlashloanReceiver.sol
+++ b/tests/Contracts/FlashloanReceiver.sol
@@ -8,9 +8,12 @@ import "../../contracts/SafeMath.sol";
 contract FlashloanReceiver is IFlashloanReceiver {
     using SafeMath for uint256;
 
+    uint totalBorrows;
+
     function doFlashloan(address cToken, uint borrowAmount, uint repayAmount) external {
         uint balanceBefore = ERC20(CCapableErc20(cToken).underlying()).balanceOf(address(this));
         bytes memory data = abi.encode(cToken, borrowAmount, repayAmount);
+        totalBorrows = CCapableErc20(cToken).totalBorrows();
         CCapableErc20(cToken).flashLoan(address(this), borrowAmount, data);
         uint balanceAfter = ERC20(CCapableErc20(cToken).underlying()).balanceOf(address(this));
         require(balanceAfter == balanceBefore.add(borrowAmount).sub(repayAmount), "Balance inconsistent");
@@ -20,6 +23,8 @@ contract FlashloanReceiver is IFlashloanReceiver {
       (address cToken, uint borrowAmount, uint repayAmount) = abi.decode(params, (address, uint, uint));
       require(underlying == CCapableErc20(cToken).underlying(), "Params not match");
       require(amount == borrowAmount, "Params not match");
+      uint totalBorrowsAfter = CCapableErc20(cToken).totalBorrows();
+      require(totalBorrows.add(borrowAmount) == totalBorrowsAfter, "totalBorrow mismatch");
       require(ERC20(underlying).transfer(cToken, repayAmount), "Transfer fund back failed");
     }
 }