Address OZ's feedback on USDT comet support

contracts/Comet.sol

@@ -201,6 +201,42 @@
         (asset14_a, asset14_b) = getPackedAssetInternal(config.assetConfigs, 14);
     }
 
+    /**
+     * @dev Prevents marked functions from being reentered 
+     */
+    modifier nonReentrant() {
+        nonReentrantBefore();
+        _;
+        nonReentrantAfter();
+    }
+
+    /**
+     * @dev Checks that the reentrancy flag is not set and then sets the flag
+     */
+    function nonReentrantBefore() internal {
+        bytes32 slot = REENTRANCY_GUARD_FLAG_SLOT;
+        uint256 status;
+        assembly ("memory-safe") {
+            status := sload(slot)
+        }
+
+        if (status == REENTRANCY_GUARD_ENTERED) revert ReentrantCallBlocked();
+        assembly ("memory-safe") {
+            sstore(slot, REENTRANCY_GUARD_ENTERED)
+        }
+    }
+
+    /**
+     * @dev Unsets the reentrancy flag
+     */
+    function nonReentrantAfter() internal {
+        bytes32 slot = REENTRANCY_GUARD_FLAG_SLOT;
+        uint256 status;
+        assembly ("memory-safe") {
+            sstore(slot, REENTRANCY_GUARD_NOT_ENTERED)
+        }
+    }
+
     /**
      * @notice Initialize storage for the contract
      * @dev Can be used from constructor or proxy
@@ -224,7 +260,7 @@
     function getPackedAssetInternal(AssetConfig[] memory assetConfigs, uint i) internal view returns (uint256, uint256) {
         AssetConfig memory assetConfig;
         if (i < assetConfigs.length) {
             assembly {
                 assetConfig := mload(add(add(assetConfigs, 0x20), mul(i, 0x20)))
             }
         } else {
@@ -767,7 +803,7 @@
         uint256 preTransferBalance = IERC20NonStandard(asset).balanceOf(address(this));
         IERC20NonStandard(asset).transferFrom(from, address(this), amount);
         bool success;
-        assembly {
+        assembly ("memory-safe") {
             switch returndatasize()
                 case 0 {                       // This is a non-standard ERC-20
                     success := not(0)          // set success to true
@@ -791,7 +827,7 @@
     function doTransferOut(address asset, address to, uint amount) internal {
         IERC20NonStandard(asset).transfer(to, amount);
         bool success;
-        assembly {
+        assembly ("memory-safe") {
             switch returndatasize()
                 case 0 {                       // This is a non-standard ERC-20
                     success := not(0)          // set success to true
@@ -841,7 +877,7 @@
      * @dev Supply either collateral or base asset, depending on the asset, if operator is allowed
      * @dev Note: Specifying an `amount` of uint256.max will repay all of `dst`'s accrued base borrow balance
      */
-    function supplyInternal(address operator, address from, address dst, address asset, uint amount) internal {
+    function supplyInternal(address operator, address from, address dst, address asset, uint amount) internal nonReentrant {
         if (isSupplyPaused()) revert Paused();
         if (!hasPermission(from, operator)) revert Unauthorized();
 
@@ -952,7 +988,7 @@
      * @dev Transfer either collateral or base asset, depending on the asset, if operator is allowed
      * @dev Note: Specifying an `amount` of uint256.max will transfer all of `src`'s accrued base balance
      */
-    function transferInternal(address operator, address src, address dst, address asset, uint amount) internal {
+    function transferInternal(address operator, address src, address dst, address asset, uint amount) internal nonReentrant {
         if (isTransferPaused()) revert Paused();
         if (!hasPermission(src, operator)) revert Unauthorized();
         if (src == dst) revert NoSelfTransfer();
@@ -1063,7 +1099,7 @@
      * @dev Withdraw either collateral or base asset, depending on the asset, if operator is allowed
      * @dev Note: Specifying an `amount` of uint256.max will withdraw all of `src`'s accrued base balance
      */
-    function withdrawInternal(address operator, address src, address to, address asset, uint amount) internal {
+    function withdrawInternal(address operator, address src, address to, address asset, uint amount) internal nonReentrant {
         if (isWithdrawPaused()) revert Paused();
         if (!hasPermission(src, operator)) revert Unauthorized();
 
@@ -1224,7 +1260,7 @@
      * @param baseAmount The amount of base tokens used to buy the collateral
      * @param recipient The recipient address
      */
-    function buyCollateral(address asset, uint minAmount, uint baseAmount, address recipient) override external {
+    function buyCollateral(address asset, uint minAmount, uint baseAmount, address recipient) override external nonReentrant {
         if (isBuyPaused()) revert Paused();
 
         int reserves = getReserves();
@@ -1286,6 +1322,7 @@
      * @dev Only callable by governor
      * @dev Note: Setting the `asset` as Comet's address will allow the manager
      * to withdraw from Comet's Comet balance
+     * @dev Note: For USDT, if there is non-zero prior allowance, it must be reset to 0 first before setting a new value in proposal
      * @param asset The asset that the manager will gain approval of
      * @param manager The account which will be allowed or disallowed
      * @param amount The amount of an asset to approve
@@ -1343,9 +1380,9 @@
     /**
      * @notice Fallback to calling the extension delegate for everything else
      */
     fallback() external payable {
         address delegate = extensionDelegate;
         assembly {
             calldatacopy(0, 0, calldatasize())
             let result := delegatecall(gas(), delegate, 0, calldatasize(), 0, 0)
             returndatacopy(0, 0, returndatasize())

contracts/CometCore.sol

@@ -56,6 +56,13 @@ abstract contract CometCore is CometConfiguration, CometStorage, CometMath {
     /// @dev The scale for factors
     uint64 internal constant FACTOR_SCALE = 1e18;
 
+    /// @dev The storage slot for reentrancy guard flags
+    bytes32 internal constant REENTRANCY_GUARD_FLAG_SLOT = bytes32(keccak256("comet.reentrancy.guard"));
+
+    /// @dev The reentrancy guard statuses
+    uint256 internal constant REENTRANCY_GUARD_NOT_ENTERED = 0;
+    uint256 internal constant REENTRANCY_GUARD_ENTERED = 1;
+
     /**
      * @notice Determine if the manager has permission to act on behalf of the owner
      * @param owner The owner account

contracts/CometMainInterface.sol

@@ -25,6 +25,7 @@ abstract contract CometMainInterface is CometCore {
     error NotForSale();
     error NotLiquidatable();
     error Paused();
+    error ReentrantCallBlocked();
     error SupplyCapExceeded();
     error TimestampTooLarge();
     error TooManyAssets();

contracts/IERC20NonStandard.sol

@@ -7,9 +7,37 @@ pragma solidity 0.8.15;
  *  See https://medium.com/coinmonks/missing-return-value-bug-at-least-130-tokens-affected-d67bf08521ca
  */
 interface IERC20NonStandard {
+    function name() external view returns (string memory);
+    function symbol() external view returns (string memory);
     function decimals() external view returns (uint8);
+
+    /**
+     * @notice Approve `spender` to transfer up to `amount` from `src`
+     * @dev This will overwrite the approval amount for `spender`
+     *  and is subject to issues noted [here](https://eips.ethereum.org/EIPS/eip-20#approve)
+     * @param spender The address of the account which may transfer tokens
+     * @param amount The number of tokens that are approved (-1 means infinite)
+     */
     function approve(address spender, uint256 amount) external;
+
+    /**
+     * @notice Transfer `value` tokens from `msg.sender` to `to`
+     * @param to The address of the destination account
+     * @param value The number of tokens to transfer
+     */
     function transfer(address to, uint256 value) external;
+
+    /**
+     * @notice Transfer `value` tokens from `from` to `to`
+     * @param from The address of the source account
+     * @param to The address of the destination account
+     * @param value The number of tokens to transfer
+     */
     function transferFrom(address from, address to, uint256 value) external;
+
+    /**
+     * @notice Gets the balance of the specified address
+     * @param account The address from which the balance will be retrieved
+     */
     function balanceOf(address account) external view returns (uint256);
-}
+}

contracts/test/EvilToken.sol

@@ -13,7 +13,8 @@ contract EvilToken is FaucetToken {
     enum AttackType {
         TRANSFER_FROM,
         WITHDRAW_FROM,
-        SUPPLY_FROM
+        SUPPLY_FROM,
+        BUY_COLLATERAL
     }
 
     struct ReentryAttack {
@@ -92,6 +93,13 @@ contract EvilToken is FaucetToken {
                 reentryAttack.asset,
                 reentryAttack.amount
             );
+        }  else if (reentryAttack.attackType == AttackType.BUY_COLLATERAL) {
+            Comet(payable(msg.sender)).buyCollateral(
+                reentryAttack.asset,
+                0,
+                reentryAttack.amount,
+                reentryAttack.destination
+            );
         } else {
             revert("invalid reentry attack");
         }

test/buy-collateral-test.ts

@@ -409,7 +409,7 @@ describe('buyCollateral', function () {
   });
 
   describe('reentrancy', function() {
-    it('is not broken by reentrancy supply ', async () => {
+    it('is not broken by reentrancy supply', async () => {
       const wethArgs = {
         initial: 1e4,
         decimals: 18,
@@ -543,8 +543,8 @@ describe('buyCollateral', function () {
       expect(normalTotalsBasic.trackingSupplyIndex).to.equal(evilTotalsBasic.trackingSupplyIndex);
       expect(normalTotalsBasic.trackingBorrowIndex).to.equal(evilTotalsBasic.trackingBorrowIndex);
       expect(normalTotalsBasic.totalSupplyBase).to.equal(1e6);
-      // EvilToken attack is using 3000 EVIL tokens, so totalSupplyBase should have 3000e6
-      expect(evilTotalsBasic.totalSupplyBase).to.equal(3000e6);
+      // EvilToken attack should be blocked
+      expect(evilTotalsBasic.totalSupplyBase).to.equal(0);
       expect(normalTotalsBasic.totalBorrowBase).to.equal(evilTotalsBasic.totalBorrowBase);
 
       expect(normalTotalsCollateral.totalSupplyAsset).to.eq(evilTotalsCollateral.totalSupplyAsset);
@@ -559,8 +559,72 @@ describe('buyCollateral', function () {
       const evilBobPortfolio = await portfolio(evilProtocol, evilBob.address);
 
       expect(normalBobPortfolio.internal.USDC).to.equal(1e6);
-      // EvilToken attack is using 3000 EVIL tokens, so totalSupplyBase should be 3000e6 (under Bob's name)
-      expect(evilBobPortfolio.internal.EVIL).to.equal(3000e6);
+      // EvilToken attack should be blocked, so totalSupplyBase should be 0
+      expect(evilBobPortfolio.internal.EVIL).to.equal(0);
+    });
+
+    it('reentrant buyCollateral is reverted', async () => {
+      const wethArgs = {
+        initial: 1e4,
+        decimals: 18,
+        initialPrice: 3000,
+      };
+      const baseTokenArgs = {
+        decimals: 6,
+        initial: 1e6,
+        initialPrice: 1,
+      };
+
+      // malicious scenario, EVIL token is base
+      const evilProtocol = await makeProtocol({
+        base: 'EVIL',
+        assets: {
+          EVIL: {
+            ...baseTokenArgs,
+            factory: await ethers.getContractFactory('EvilToken') as EvilToken__factory,
+          },
+          WETH: wethArgs,
+        },
+        targetReserves: 1
+      });
+      const {
+        comet: evilComet,
+        tokens: evilTokens,
+        users: [evilAlice, evilBob]
+      } = evilProtocol;
+      const { WETH: evilWETH, EVIL } = <{ WETH: FaucetToken, EVIL: EvilToken }>evilTokens;
+
+      // add attack to EVIL token
+      const attack = Object.assign({}, await EVIL.getAttack(), {
+        attackType: ReentryAttack.BuyCollateral,
+        source: evilAlice.address,
+        destination: evilBob.address,
+        asset: evilWETH.address,
+        amount: 3000e6,
+        maxCalls: 1
+      });
+      await EVIL.setAttack(attack);
+
+      // allocate tokens (evil)
+      await evilWETH.allocateTo(evilComet.address, exp(100, 18));
+      await EVIL.allocateTo(evilAlice.address, exp(5000, 6));
+
+      // approve Comet to move funds
+      await EVIL.connect(evilAlice).approve(EVIL.address, exp(5000, 6));
+      await EVIL.connect(evilAlice).approve(evilComet.address, exp(5000, 6));
+
+      // authorize EVIL, since callback will originate from EVIL token address
+      await evilComet.connect(evilAlice).allow(EVIL.address, true);
+
+      // call buyCollateral; supplyFrom is called in callback
+      await expect(evilComet
+        .connect(evilAlice)
+        .buyCollateral(
+          evilWETH.address,
+          exp(0, 18),
+          exp(3000, 6),
+          evilAlice.address
+        )).to.be.revertedWith("custom error 'ReentrantCallBlocked()'");
     });
   });
 });

test/helpers.ts

@@ -44,7 +44,8 @@ export type Numeric = number | bigint;
 export enum ReentryAttack {
   TransferFrom = 0,
   WithdrawFrom = 1,
-  SupplyFrom = 2
+  SupplyFrom = 2,
+  BuyCollateral = 3,
 }
 
 export type ProtocolOpts = {

test/supply-test.ts

@@ -52,7 +52,7 @@ describe('supplyTo', function () {
     expect(q1.external).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n });
     expect(t1.totalSupplyBase).to.be.equal(t0.totalSupplyBase.add(100e6));
     expect(t1.totalBorrowBase).to.be.equal(t0.totalBorrowBase);
-    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(120000);
+    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(124000);
   });
 
   it('supplies max base borrow balance (including accrued) from sender if the asset is base', async () => {
@@ -259,7 +259,7 @@ describe('supplyTo', function () {
     expect(p1.external).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n });
     expect(t1.totalSupplyBase).to.be.equal(109);
     expect(t1.totalBorrowBase).to.be.equal(t0.totalBorrowBase);
-    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(120000);
+    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(124000);
   });
 
   it('supplies collateral from sender if the asset is collateral', async () => {
@@ -305,7 +305,7 @@ describe('supplyTo', function () {
     expect(q1.internal).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n });
     expect(q1.external).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n });
     expect(t1.totalSupplyAsset).to.be.equal(t0.totalSupplyAsset.add(8e8));
-    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(140000);
+    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(153000);
   });
 
   it('calculates base principal correctly', async () => {
@@ -466,7 +466,7 @@ describe('supplyTo', function () {
     expect(t1.totalSupplyBase).to.be.equal(t0.totalSupplyBase.add(999e6));
     expect(t1.totalBorrowBase).to.be.equal(t0.totalBorrowBase);
     // Fee Token logics will cost a bit more gas than standard ERC20 token with no fee calculation
-    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(128000);
+    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(151000);
   });
 
   it('supplies collateral the correct amount in a fee-like situation', async () => {
@@ -524,10 +524,10 @@ describe('supplyTo', function () {
     expect(q1.external).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n, FeeToken: 0n });
     expect(t1.totalSupplyAsset).to.be.equal(t0.totalSupplyAsset.add(1998e8));
     // Fee Token logics will cost a bit more gas than standard ERC20 token with no fee calculation
-    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(163000);
+    expect(Number(s0.receipt.gasUsed)).to.be.lessThan(186000);
   });
 
-  it('prevents exceeding the supply cap via re-entrancy', async () => {
+  it('blocks reentrancy from exceeding the supply cap', async () => {
     const { comet, tokens, users: [alice, bob] } = await makeProtocol({
       assets: {
         USDC: {
@@ -558,47 +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("custom error 'SupplyCapExceeded()'");
-  });
-
-  // This test is purposely test the edge case where comet internal accounting will be incorrect when EvilToken re-entrancy attack
-  // is performed. The attack will cause comet to have inflated supply amount of token. 
-  //
-  // We decided that this won't be a concern for Comet, as in order to pull off this attack, the governance has to propose and vote to 
-  // add suspicious token in to Comet's market. As long as governance doesn't add suspicious token contract or erc-777 token to market, the 
-  // Comet should not be vulnerable to this type of attack.
-  it('incorrect accounting via re-entrancy', async () => {
-    const { comet, tokens, users: [alice, bob] } = await makeProtocol({
-      assets: {
-        USDC: {
-          decimals: 6
-        },
-        EVIL: {
-          decimals: 6,
-          initialPrice: 2,
-          factory: await ethers.getContractFactory('EvilToken') as EvilToken__factory,
-          supplyCap: 250e6
-        }
-      }
-    });
-    const { EVIL } = <{ EVIL: EvilToken }>tokens;
-
-    const attack = Object.assign({}, await EVIL.getAttack(), {
-      attackType: ReentryAttack.SupplyFrom,
-      source: alice.address,
-      destination: bob.address,
-      asset: EVIL.address,
-      amount: 75e6,
-      maxCalls: 1
-    });
-    await EVIL.setAttack(attack);
-
-    await comet.connect(alice).allow(EVIL.address, true);
-    await wait(EVIL.connect(alice).approve(comet.address, 75e6));
-    await EVIL.allocateTo(alice.address, 75e6);
-    await comet.connect(alice).supplyTo(bob.address, EVIL.address, 75e6);
-    // Re-entrancy attack 2 loops, expect to have 2X accounting balance 75*2 = 150
-    expect(await comet.collateralBalanceOf(bob.address, EVIL.address)).to.be.equal(150e6);
+    ).to.be.revertedWithCustomError(comet, 'ReentrantCallBlocked');
   });
 });