Skip to content

Commit

Permalink
add reentrancy guard, and updated tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Hans Wang committed Nov 8, 2023
1 parent 3544043 commit 85be018
Show file tree
Hide file tree
Showing 7 changed files with 64 additions and 70 deletions.
30 changes: 28 additions & 2 deletions contracts/Comet.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,32 @@ contract Comet is CometMainInterface {
(asset14_a, asset14_b) = getPackedAssetInternal(config.assetConfigs, 14);
}

modifier nonReentrant() {
_nonReentrantBefore();
_;
_nonReentrantAfter();
}

function _nonReentrantBefore() private {
bytes32 slot = COMET_REENTRANCY_GUARD_FLAG_SLOT;
uint256 status;
assembly {

Check warning on line 213 in contracts/Comet.sol

View workflow job for this annotation

GitHub Actions / Contract linter 

Avoid to use inline assembly. It is acceptable only in rare cases
status := sload(slot)
}

if(status == COMET_REENTRANCY_GUARD_ENTERED) revert ReentrantCallBlocked();
assembly {

Check warning on line 218 in contracts/Comet.sol

View workflow job for this annotation

GitHub Actions / Contract linter 

Avoid to use inline assembly. It is acceptable only in rare cases
sstore(slot, COMET_REENTRANCY_GUARD_ENTERED)
}
}

function _nonReentrantAfter() private {
bytes32 slot = COMET_REENTRANCY_GUARD_FLAG_SLOT;
assembly {

Check warning on line 225 in contracts/Comet.sol

View workflow job for this annotation

GitHub Actions / Contract linter 

Avoid to use inline assembly. It is acceptable only in rare cases
sstore(slot, COMET_REENTRANCY_GUARD_NOT_ENTERED)
}
}

/**
* @notice Initialize storage for the contract
* @dev Can be used from constructor or proxy
Expand Down Expand Up @@ -763,7 +789,7 @@ contract Comet is CometMainInterface {
* @dev Safe ERC20 transfer in and returns the final amount transferred (taking into account any fees)
* @dev Note: Safely handles non-standard ERC-20 tokens that do not return a value. See here: https://medium.com/coinmonks/missing-return-value-bug-at-least-130-tokens-affected-d67bf08521ca
*/
function doTransferIn(address asset, address from, uint amount) internal returns (uint) {
function doTransferIn(address asset, address from, uint amount) nonReentrant internal returns (uint) {

Check warning on line 792 in contracts/Comet.sol

View workflow job for this annotation

GitHub Actions / Contract linter 

Visibility modifier must be first in list of modifiers
uint256 preTransferBalance = ERC20(asset).balanceOf(address(this));
(bool success, bytes memory returndata) = asset.call(abi.encodeWithSelector(ERC20.transferFrom.selector, from, address(this), amount));

Check warning on line 794 in contracts/Comet.sol

View workflow job for this annotation

GitHub Actions / Contract linter 

Avoid to use low level calls
if (!success || !(returndata.length == 0 || abi.decode(returndata, (bool)))) {
Expand All @@ -776,7 +802,7 @@ contract Comet is CometMainInterface {
* @dev Safe ERC20 transfer out
* @dev Note: Safely handles non-standard ERC-20 tokens that do not return a value. See here: https://medium.com/coinmonks/missing-return-value-bug-at-least-130-tokens-affected-d67bf08521ca
*/
function doTransferOut(address asset, address to, uint amount) internal {
function doTransferOut(address asset, address to, uint amount) nonReentrant internal {

Check warning on line 805 in contracts/Comet.sol

View workflow job for this annotation

GitHub Actions / Contract linter 

Visibility modifier must be first in list of modifiers
(bool success, bytes memory returndata) = asset.call(abi.encodeWithSelector(ERC20.transfer.selector, to, amount));

Check warning on line 806 in contracts/Comet.sol

View workflow job for this annotation

GitHub Actions / Contract linter 

Avoid to use low level calls
if (!success || !(returndata.length == 0 || abi.decode(returndata, (bool)))) {
revert TransferOutFailed();
Expand Down
7 changes: 7 additions & 0 deletions contracts/CometCore.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 COMET_REENTRANCY_GUARD_FLAG_SLOT = bytes32(keccak256("comet.reentrancy.guard"));

/// @dev The reentrancy guard status
uint256 internal constant COMET_REENTRANCY_GUARD_NOT_ENTERED = 1;
uint256 internal constant COMET_REENTRANCY_GUARD_ENTERED = 2;

/**
* @notice Determine if the manager has permission to act on behalf of the owner
* @param owner The owner account
Expand Down
1 change: 1 addition & 0 deletions contracts/CometMainInterface.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ abstract contract CometMainInterface is CometCore {
error TransferInFailed();
error TransferOutFailed();
error Unauthorized();
error ReentrantCallBlocked();

event Supply(address indexed from, address indexed dst, uint amount);
event Transfer(address indexed from, address indexed to, uint amount);
Expand Down
16 changes: 8 additions & 8 deletions contracts/IERC20NonStandard.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,23 +21,23 @@ interface IERC20NonStandard {
function approve(address spender, uint256 amount) external;

/**
* @notice Transfer `amount` tokens from `msg.sender` to `dst`
* @param dst The address of the destination account
* @param amount The number of tokens to transfer
* @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 `amount` tokens from `src` to `dst`
* @param src The address of the source account
* @param dst The address of the destination account
* @param amount The number of tokens to transfer
* @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 owner The address from which the balance will be retrieved
* @param account The address from which the balance will be retrieved
*/
function balanceOf(address account) external view returns (uint256);
}
10 changes: 5 additions & 5 deletions test/buy-collateral-test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ describe('buyCollateral', function () {
});

describe('reentrancy', function() {
it('is not broken by reentrancy supply ', async () => {
it('block reentrancy supply ', async () => {
const wethArgs = {
initial: 1e4,
decimals: 18,
Expand Down Expand Up @@ -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);
Expand All @@ -559,8 +559,8 @@ 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);
});
});
});
56 changes: 8 additions & 48 deletions test/supply-test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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('re-entrancy block', async () => {
const { comet, tokens, users: [alice, bob] } = await makeProtocol({
assets: {
USDC: {
Expand All @@ -537,7 +537,7 @@ describe('supplyTo', function () {
decimals: 6,
initialPrice: 2,
factory: await ethers.getContractFactory('EvilToken') as EvilToken__factory,
supplyCap: 100e6
supplyCap: 250e6
}
}
});
Expand All @@ -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.revertedWith("custom error 'TransferInFailed()'");
});
});

Expand Down
14 changes: 7 additions & 7 deletions test/withdraw-test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ describe('withdrawTo', function () {
expect(q1.external).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n });
expect(t1.totalSupplyBase).to.be.equal(0n);
expect(t1.totalBorrowBase).to.be.equal(0n);
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(100000);
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(106000);
});

it('does not emit Transfer for 0 burn', async () => {
Expand Down Expand Up @@ -144,7 +144,7 @@ describe('withdrawTo', function () {
expect(b1.external).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n });
expect(t1.totalSupplyBase).to.be.equal(0n);
expect(t1.totalBorrowBase).to.be.equal(exp(50, 6));
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(110000);
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(115000);
});

it('withdraw max base should withdraw 0 if user has a borrow position', async () => {
Expand Down Expand Up @@ -190,7 +190,7 @@ describe('withdrawTo', function () {
expect(b1.external).to.be.deep.equal({ USDC: 0n, COMP: 0n, WETH: 0n, WBTC: 0n });
expect(t1.totalSupplyBase).to.be.equal(t0.totalSupplyBase);
expect(t1.totalBorrowBase).to.be.equal(t0.totalBorrowBase);
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(110000);
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(121000);
});

// This demonstrates a weird quirk of the present value/principal value rounding down math.
Expand Down Expand Up @@ -279,7 +279,7 @@ describe('withdrawTo', 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(0n);
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(80000);
expect(Number(s0.receipt.gasUsed)).to.be.lessThan(85000);
});

it('calculates base principal correctly', async () => {
Expand Down Expand Up @@ -475,7 +475,7 @@ describe('withdraw', function () {
});

describe('reentrancy', function () {
it('is not broken by malicious reentrancy transferFrom', async () => {
it('block malicious reentrancy transferFrom', async () => {
const { comet, tokens, users: [alice, bob] } = await makeProtocol({
assets: {
USDC: {
Expand Down Expand Up @@ -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("custom error 'NotCollateralized()'");
).to.be.revertedWith("custom error 'TransferOutFailed()'");

// no USDC transferred
expect(await USDC.balanceOf(comet.address)).to.eq(100e6);
Expand Down Expand Up @@ -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("custom error 'NotCollateralized()'");
).to.be.revertedWith("custom error 'TransferOutFailed()'");

// no USDC transferred
expect(await USDC.balanceOf(comet.address)).to.eq(100e6);
Expand Down

0 comments on commit 85be018

Please sign in to comment.