reentrancy guard in Comet, isntead of importing external contracts wi…

…th the risk of storage collision
vincetiu8 · Sep 16, 2023 · d0f016a · d0f016a
1 parent 49588d6
commit d0f016a
Show file tree

Hide file tree

Showing 7 changed files with 36 additions and 84 deletions.
diff --git a/contracts/Comet.sol b/contracts/Comet.sol
@@ -4,14 +4,13 @@ pragma solidity 0.8.15;
 import "./CometMainInterface.sol";
 import "./IERC20NonStandard.sol";
 import "./IPriceFeed.sol";
-import "./ReentrancyGuard.sol";
 
 /**
  * @title Compound's Comet Contract
  * @notice An efficient monolithic money market protocol
  * @author Compound
  */
-contract Comet is CometMainInterface, ReentrancyGuard {
+contract Comet is CometMainInterface {
     /** General configuration constants **/
 
     /// @notice The admin of the protocol
@@ -213,7 +212,7 @@ contract Comet is CometMainInterface, ReentrancyGuard {
         lastAccrualTime = getNowInternal();
         baseSupplyIndex = BASE_INDEX_SCALE;
         baseBorrowIndex = BASE_INDEX_SCALE;
-
+        reentrancyGuardStatus = REENTRANCY_GUARD_MUTEX_NOT_ENTERED;
         // Implicit initialization (not worth increasing contract size)
         // trackingSupplyIndex = 0;
         // trackingBorrowIndex = 0;
@@ -1341,6 +1340,28 @@ contract Comet is CometMainInterface, ReentrancyGuard {
         return principal < 0 ? presentValueBorrow(baseBorrowIndex_, unsigned104(-principal)) : 0;
     }
 
+    /**
+     * @notice Reentrancy guard function and modifier for nonReentrant functions    
+     * @dev Note: reference: https://github.com/OpenZeppelin/openzeppelin-contracts/blob/v4.9.3/contracts/security/ReentrancyGuard.sol
+     */
+    modifier nonReentrant() {
+        _nonReentrantBefore();
+        _;
+        _nonReentrantAfter();
+    }
+
+    function _nonReentrantBefore() private {
+        if (reentrancyGuardStatus == REENTRANCY_GUARD_MUTEX_ENTERED) {
+            revert ReentrancyGuardReentrantCall();
+        }
+
+        reentrancyGuardStatus = REENTRANCY_GUARD_MUTEX_ENTERED;
+    }
+
+    function _nonReentrantAfter() private {
+        reentrancyGuardStatus = REENTRANCY_GUARD_MUTEX_NOT_ENTERED;
+    }
+
     /**
      * @notice Fallback to calling the extension delegate for everything else
      */

diff --git a/contracts/CometCore.sol b/contracts/CometCore.sol
@@ -56,6 +56,10 @@ abstract contract CometCore is CometConfiguration, CometStorage, CometMath {
     /// @dev The scale for factors
     uint64 internal constant FACTOR_SCALE = 1e18;
 
+    /// @dev Reentrancy guard mutex status
+    uint256 internal constant REENTRANCY_GUARD_MUTEX_NOT_ENTERED = 1;
+    uint256 internal constant REENTRANCY_GUARD_MUTEX_ENTERED = 2;
+
     /**
      * @notice Determine if the manager has permission to act on behalf of the owner
      * @param owner The owner account

diff --git a/contracts/CometMainInterface.sol b/contracts/CometMainInterface.sol
@@ -32,6 +32,7 @@ abstract contract CometMainInterface is CometCore {
     error TransferInFailed();
     error TransferOutFailed();
     error Unauthorized();
+    error ReentrancyGuardReentrantCall();
 
     event Supply(address indexed from, address indexed dst, uint amount);
     event Transfer(address indexed from, address indexed to, uint amount);
@@ -55,7 +56,7 @@ abstract contract CometMainInterface is CometCore {
 
     /// @notice Event emitted when reserves are withdrawn by the governor
     event WithdrawReserves(address indexed to, uint amount);
-
+    
     function supply(address asset, uint amount) virtual external;
     function supplyTo(address dst, address asset, uint amount) virtual external;
     function supplyFrom(address from, address dst, address asset, uint amount) virtual external;

diff --git a/contracts/CometStorage.sol b/contracts/CometStorage.sol
@@ -73,4 +73,7 @@ contract CometStorage {
 
     /// @notice Mapping of magic liquidator points
     mapping(address => LiquidatorPoints) public liquidatorPoints;
+
+    /// @dev Reentrancy guard mutex status
+    uint256 internal reentrancyGuardStatus;
 }
diff --git a/contracts/ReentrancyGuard.sol b/contracts/ReentrancyGuard.sol
diff --git a/test/supply-test.ts b/test/supply-test.ts
@@ -558,7 +558,7 @@ describe('supplyTo', function () {
     await EVIL.allocateTo(alice.address, 75e6);
     await expect(
       comet.connect(alice).supplyTo(bob.address, EVIL.address, 75e6)
-    ).to.be.revertedWith("ReentrancyGuard: reentrant call");
+    ).to.be.revertedWith("custom error 'ReentrancyGuardReentrantCall()'");
   });
 });
 

diff --git a/test/withdraw-test.ts b/test/withdraw-test.ts
@@ -511,7 +511,7 @@ describe('withdraw', function () {
       // in callback, EVIL token calls transferFrom(alice.address, bob.address, 1e6)
       await expect(
         comet.connect(alice).withdraw(EVIL.address, 1e6)
-      ).to.be.revertedWith("ReentrancyGuard: reentrant call");
+      ).to.be.revertedWith("custom error 'ReentrancyGuardReentrantCall()'");
 
       // no USDC transferred
       expect(await USDC.balanceOf(comet.address)).to.eq(100e6);
@@ -558,7 +558,7 @@ describe('withdraw', function () {
       // in callback, EvilToken attempts to withdraw USDC to bob's address
       await expect(
         comet.connect(alice).withdraw(EVIL.address, 1e6)
-      ).to.be.revertedWith("ReentrancyGuard: reentrant call");
+      ).to.be.revertedWith("custom error 'ReentrancyGuardReentrantCall()'");
 
       // no USDC transferred
       expect(await USDC.balanceOf(comet.address)).to.eq(100e6);