Add internal overrideable _flashFee in ERC20FlashMint (#3551)

Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com> Co-authored-by: Francisco <frangio.1@gmail.com>
OpenZeppelin · Jul 22, 2022 · d1b1e17 · d1b1e17
1 parent 6e8d885
commit d1b1e17
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@
  * `ERC721`: optimize burn by making approval clearing implicit instead of emitting an event. ([#3538](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3538))
  * `ReentrancyGuard`: Reduce code size impact of the modifier by using internal functions. ([#3515](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3515))
  * `SafeCast`: optimize downcasting of signed integers. ([#3565](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3565))
+ * `ERC20FlashMint`: add an internal `_flashFee` function for overriding. ([#3551](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3551))
 
 ### Compatibility Note
 

diff --git a/contracts/mocks/ERC20FlashMintMock.sol b/contracts/mocks/ERC20FlashMintMock.sol
@@ -25,8 +25,7 @@ contract ERC20FlashMintMock is ERC20FlashMint {
         _flashFeeAmount = amount;
     }
 
-    function flashFee(address token, uint256 amount) public view virtual override returns (uint256) {
-        super.flashFee(token, amount);
+    function _flashFee(address, uint256) internal view override returns (uint256) {
         return _flashFeeAmount;
     }
 

diff --git a/contracts/token/ERC20/extensions/ERC20FlashMint.sol b/contracts/token/ERC20/extensions/ERC20FlashMint.sol
@@ -38,7 +38,18 @@ abstract contract ERC20FlashMint is ERC20, IERC3156FlashLender {
      */
     function flashFee(address token, uint256 amount) public view virtual override returns (uint256) {
         require(token == address(this), "ERC20FlashMint: wrong token");
+        return _flashFee(token, amount);
+    }
+
+    /**
+     * @dev Returns the fee applied when doing flash loans. This function calls the {flashFee} function which returns the fee applied when doing flash loans.
+     * @param token The token to be flash loaned.
+     * @param amount The amount of tokens to be loaned.
+     * @return The fees applied to the corresponding flash loan.
+     */
+    function _flashFee(address token, uint256 amount) internal view virtual returns (uint256) {
         // silence warning about unused variable without the addition of bytecode.
+        token;
         amount;
         return 0;
     }

diff --git a/test/token/ERC20/extensions/ERC20FlashMint.test.js b/test/token/ERC20/extensions/ERC20FlashMint.test.js
@@ -8,7 +8,7 @@ const ERC20FlashMintMock = artifacts.require('ERC20FlashMintMock');
 const ERC3156FlashBorrowerMock = artifacts.require('ERC3156FlashBorrowerMock');
 
 contract('ERC20FlashMint', function (accounts) {
-  const [ initialHolder, other, anotherAccount ] = accounts;
+  const [initialHolder, other, anotherAccount] = accounts;
 
   const name = 'My Token';
   const symbol = 'MTKN';
@@ -61,23 +61,23 @@ contract('ERC20FlashMint', function (accounts) {
       expect(await this.token.allowance(receiver.address, this.token.address)).to.be.bignumber.equal('0');
     });
 
-    it ('missing return value', async function () {
+    it('missing return value', async function () {
       const receiver = await ERC3156FlashBorrowerMock.new(false, true);
       await expectRevert(
         this.token.flashLoan(receiver.address, this.token.address, loanAmount, '0x'),
         'ERC20FlashMint: invalid return value',
       );
     });
 
-    it ('missing approval', async function () {
+    it('missing approval', async function () {
       const receiver = await ERC3156FlashBorrowerMock.new(true, false);
       await expectRevert(
         this.token.flashLoan(receiver.address, this.token.address, loanAmount, '0x'),
         'ERC20: insufficient allowance',
       );
     });
 
-    it ('unavailable funds', async function () {
+    it('unavailable funds', async function () {
       const receiver = await ERC3156FlashBorrowerMock.new(true, true);
       const data = this.token.contract.methods.transfer(other, 10).encodeABI();
       await expectRevert(
@@ -86,7 +86,7 @@ contract('ERC20FlashMint', function (accounts) {
       );
     });
 
-    it ('more than maxFlashLoan', async function () {
+    it('more than 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.
@@ -96,8 +96,8 @@ contract('ERC20FlashMint', function (accounts) {
     describe('custom flash fee & custom fee receiver', function () {
       const receiverInitialBalance = new BN(200000);
       const flashFee = new BN(5000);
-      
-      beforeEach('init reciever balance & set flash fee',async function () {        
+
+      beforeEach('init reciever balance & set flash fee', async function () {
         this.receiver = await ERC3156FlashBorrowerMock.new(true, true);
         const receipt = await this.token.mint(this.receiver.address, receiverInitialBalance);
         await expectEvent(receipt, 'Transfer', { from: ZERO_ADDRESS, to: this.receiver.address, value: receiverInitialBalance });
@@ -106,13 +106,13 @@ contract('ERC20FlashMint', function (accounts) {
         await this.token.setFlashFee(flashFee);
         expect(await this.token.flashFee(this.token.address, loanAmount)).to.be.bignumber.equal(flashFee);
       });
-      
+
       it('default flash fee receiver', async function () {
         const { tx } = await this.token.flashLoan(this.receiver.address, this.token.address, loanAmount, '0x');
         await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: ZERO_ADDRESS, to: this.receiver.address, value: loanAmount });
-        await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: this.receiver.address, to: ZERO_ADDRESS, value: loanAmount.add (flashFee)});
-        await expectEvent.inTransaction(tx, this.receiver, 'BalanceOf', { token: this.token.address, account: this.receiver.address, value:   receiverInitialBalance.add(loanAmount) });
-        await expectEvent.inTransaction(tx, this.receiver, 'TotalSupply', { token: this.token.address, value: initialSupply.add  (receiverInitialBalance).add(loanAmount) });
+        await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: this.receiver.address, to: ZERO_ADDRESS, value: loanAmount.add(flashFee) });
+        await expectEvent.inTransaction(tx, this.receiver, 'BalanceOf', { token: this.token.address, account: this.receiver.address, value: receiverInitialBalance.add(loanAmount) });
+        await expectEvent.inTransaction(tx, this.receiver, 'TotalSupply', { token: this.token.address, value: initialSupply.add(receiverInitialBalance).add(loanAmount) });
 
         expect(await this.token.totalSupply()).to.be.bignumber.equal(initialSupply.add(receiverInitialBalance).sub(flashFee));
         expect(await this.token.balanceOf(this.receiver.address)).to.be.bignumber.equal(receiverInitialBalance.sub(flashFee));
@@ -124,15 +124,15 @@ contract('ERC20FlashMint', function (accounts) {
         const flashFeeReceiverAddress = anotherAccount;
         await this.token.setFlashFeeReceiver(flashFeeReceiverAddress);
         expect(await this.token.flashFeeReceiver()).to.be.eq(flashFeeReceiverAddress);
-        
+
         expect(await this.token.balanceOf(flashFeeReceiverAddress)).to.be.bignumber.equal('0');
-        
+
         const { tx } = await this.token.flashLoan(this.receiver.address, this.token.address, loanAmount, '0x');
         await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: ZERO_ADDRESS, to: this.receiver.address, value: loanAmount });
         await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: this.receiver.address, to: ZERO_ADDRESS, value: loanAmount });
-        await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: this.receiver.address, to: flashFeeReceiverAddress, value: flashFee  });
-        await expectEvent.inTransaction(tx, this.receiver, 'BalanceOf', { token: this.token.address, account: this.receiver.address, value:   receiverInitialBalance.add(loanAmount) });
-        await expectEvent.inTransaction(tx, this.receiver, 'TotalSupply', { token: this.token.address, value: initialSupply.add  (receiverInitialBalance).add(loanAmount) });
+        await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: this.receiver.address, to: flashFeeReceiverAddress, value: flashFee });
+        await expectEvent.inTransaction(tx, this.receiver, 'BalanceOf', { token: this.token.address, account: this.receiver.address, value: receiverInitialBalance.add(loanAmount) });
+        await expectEvent.inTransaction(tx, this.receiver, 'TotalSupply', { token: this.token.address, value: initialSupply.add(receiverInitialBalance).add(loanAmount) });
 
         expect(await this.token.totalSupply()).to.be.bignumber.equal(initialSupply.add(receiverInitialBalance));
         expect(await this.token.balanceOf(this.receiver.address)).to.be.bignumber.equal(receiverInitialBalance.sub(flashFee));