address PR comments

OpenZeppelin · Mar 23, 2021 · b1c1f6c · b1c1f6c
1 parent 90c774d
commit b1c1f6c
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/...token/ERC20/extensions/draft-IERC3156.sol → contracts/interfaces/IERC3156.sol b/...token/ERC20/extensions/draft-IERC3156.sol → contracts/interfaces/IERC3156.sol
diff --git a/contracts/mocks/ERC3156FlashBorrowerMock.sol b/contracts/mocks/ERC3156FlashBorrowerMock.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.0;
 
 
 import "../token/ERC20/IERC20.sol";
-import "../token/ERC20/extensions/draft-IERC3156.sol";
+import "../interfaces/IERC3156.sol";
 import "../utils/Address.sol";
 
 contract ERC3156FlashBorrowerMock is IERC3156FlashBorrower {
@@ -28,6 +28,8 @@ contract ERC3156FlashBorrowerMock is IERC3156FlashBorrower {
         uint256 fee,
         bytes calldata data
     ) public override returns (bytes32) {
+        require(msg.sender == token);
+
         emit BalanceOf(token, address(this), IERC20(token).balanceOf(address(this)));
         emit TotalSupply(token, IERC20(token).totalSupply());
 

diff --git a/contracts/token/ERC20/extensions/draft-ERC20FlashMint.sol b/contracts/token/ERC20/extensions/draft-ERC20FlashMint.sol
@@ -2,7 +2,7 @@
 
 pragma solidity ^0.8.0;
 
-import "./draft-IERC3156.sol";
+import "../../../interfaces/IERC3156.sol";
 import "../ERC20.sol";
 
 /**
@@ -62,15 +62,11 @@ abstract contract ERC20FlashMint is ERC20, IERC3156FlashLender {
     {
         require(token == address(this), "ERC20FlashMint: wrong token");
         uint256 fee = flashFee(token, amount);
-        // mint tokens - will revert on overflow
         _mint(address(receiver), amount);
-        // call the flashLoan borrower
         require(receiver.onFlashLoan(msg.sender, token, amount, fee, data) == RETURN_VALUE, "ERC20FlashMint: invalid return value");
-        // update approval (equivalent of burnFrom #1) - will revert on overflow
         uint256 currentAllowance = allowance(address(receiver), address(this));
         require(currentAllowance >= amount + fee, "ERC20FlashMint: allowance does not allow refund");
         _approve(address(receiver), address(this), currentAllowance - amount - fee);
-        // burn tokens (equivalent of burnFrom #2)
         _burn(address(receiver), amount + fee);
         return true;
     }

diff --git a/test/token/ERC20/extensions/draft-ERC20FlashMint.test.js b/test/token/ERC20/extensions/draft-ERC20FlashMint.test.js
@@ -79,5 +79,12 @@ contract('ERC20FlashMint', function (accounts) {
         'ERC20: burn amount exceeds balance',
       );
     });
+
+    it ('more then maxFlashLoan', async function () {
+      const receiver = await ERC3156FlashBorrowerMock.new(true, true);
+      const data = this.token.contract.methods.transfer(other, 10).encodeABI();
+      // _mint overflow reverts using a panic code. No reason string.
+      await expectRevert.unspecified(this.token.flashLoan(receiver.address, this.token.address, MAX_UINT256, data));
+    });
   });
 });