Add customizable fee receiver to ERC20FlashMint

contracts/mocks/ERC20FlashMintFeeReceiverMock.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../token/ERC20/extensions/ERC20FlashMint.sol";
+
+contract ERC20FlashMintFeeReceiverMock is ERC20FlashMint {
+    constructor(
+        string memory name,
+        string memory symbol,
+        address initialAccount,
+        uint256 initialBalance
+    ) ERC20(name, symbol) {
+        _mint(initialAccount, initialBalance);
+    }
+
+    function mint(address account, uint256 amount) public {
+        _mint(account, amount);
+    }
+
+    function flashFee(address token, uint256 amount) public view virtual override returns (uint256) {
+        token;
+        amount;
+        return 5000;
+    }
+
+    function mockFlashFeeReceiver(address token) public view returns (address) {
+        return _flashFeeReceiver(token);
+    }
+
+    function _flashFeeReceiver(address token) internal view override returns (address) {
+        token;
+        return address(this);
+    }
+}

contracts/mocks/ERC20FlashMintMock.sol

@@ -13,4 +13,8 @@ contract ERC20FlashMintMock is ERC20FlashMint {
     ) ERC20(name, symbol) {
         _mint(initialAccount, initialBalance);
     }
+
+    function mockFlashFeeReceiver(address token) public view returns (address) {
+        return _flashFeeReceiver(token);
+    }
 }

contracts/token/ERC20/extensions/ERC20FlashMint.sol

@@ -43,6 +43,19 @@ abstract contract ERC20FlashMint is ERC20, IERC3156FlashLender {
         return 0;
     }
 
+    /**
+     * @dev Returns the reciever address of the flash fee. By default this
+     * implementation returns the address(0) which means the fee amount will be burnt.
+     * This function can be overloaded to change the fee reciever.
+     * @param token The token to be flash loaned.
+     * @return The address for which the flash fee will be sent to.
+     */
+    function _flashFeeReceiver(address token) internal view virtual returns (address) {
+        // silence warning about unused variable without the addition of bytecode.
+        token;
+        return address(0);
+    }
+
     /**
      * @dev Performs a flash loan. New tokens are minted and sent to the
      * `receiver`, who is required to implement the {IERC3156FlashBorrower}
@@ -73,8 +86,14 @@ abstract contract ERC20FlashMint is ERC20, IERC3156FlashLender {
             receiver.onFlashLoan(msg.sender, token, amount, fee, data) == _RETURN_VALUE,
             "ERC20FlashMint: invalid return value"
         );
+        address flashFeeReceiver = _flashFeeReceiver(token);
         _spendAllowance(address(receiver), address(this), amount + fee);
-        _burn(address(receiver), amount + fee);
+        if (fee == 0 || flashFeeReceiver == address(0)) {
+            _burn(address(receiver), amount + fee);
+        } else {
+            _burn(address(receiver), amount);
+            _transfer(address(receiver), flashFeeReceiver, fee);
+        }
         return true;
     }
 }

test/token/ERC20/extensions/ERC20FlashMint.test.js

@@ -40,6 +40,12 @@ contract('ERC20FlashMint', function (accounts) {
     });
   });
 
+  describe('mockFlashFeeReceiver', function () {
+    it('default receiver', async function () {
+      expect(await this.token.mockFlashFeeReceiver(this.token.address)).to.be.eq(ZERO_ADDRESS);
+    });
+  });
+
   describe('flashLoan', function () {
     it('success', async function () {
       const receiver = await ERC3156FlashBorrowerMock.new(true, true);

test/token/ERC20/extensions/ERC20FlashMintFeeReceiver.test.js

@@ -0,0 +1,61 @@
+/* eslint-disable */
+
+const { BN, constants, expectEvent, expectRevert, time } = require('@openzeppelin/test-helpers');
+const { expect } = require('chai');
+const { MAX_UINT256, ZERO_ADDRESS, ZERO_BYTES32 } = constants;
+
+const ERC20FlashMintFeeReceiverMock = artifacts.require('ERC20FlashMintFeeReceiverMock');
+const ERC3156FlashBorrowerMock = artifacts.require('ERC3156FlashBorrowerMock');
+
+contract('ERC20FlashMintFeeReceiver', function (accounts) {
+  const [ initialHolder, other ] = accounts;
+
+  const name = 'My Token';
+  const symbol = 'MTKN';
+
+  const initialSupply = new BN(100);
+  const loanAmount = new BN(10000000000000);
+
+  const flashFee = new BN(5000);
+
+  beforeEach(async function () {
+    this.token = await ERC20FlashMintFeeReceiverMock.new(name, symbol, initialHolder, initialSupply);
+  });
+
+  describe('flashFee', function () {
+    it('token amount', async function () {
+      expect(await this.token.flashFee(this.token.address, loanAmount)).to.be.bignumber.equal(flashFee);
+    });
+  });
+
+  describe('mockFlashFeeReceiver', function () {
+    it('fee receiver', async function () {
+      expect(await this.token.mockFlashFeeReceiver(this.token.address)).to.be.eq(this.token.address);
+    });
+  });
+
+  describe('flashLoan', function () {
+    it('transfered fee', async function () {
+      const receiver = await ERC3156FlashBorrowerMock.new(true, true);
+      const receiverInitialBalance = new BN(200000);
+
+      const receipt = await this.token.mint(receiver.address, receiverInitialBalance);
+      await expectEvent(receipt, 'Transfer', { from: ZERO_ADDRESS, to: receiver.address, value: receiverInitialBalance });
+
+      expect(await this.token.balanceOf(this.token.address)).to.be.bignumber.equal('0');
+      expect(await this.token.balanceOf(receiver.address)).to.be.bignumber.equal(receiverInitialBalance);
+
+      const { tx } = await this.token.flashLoan(receiver.address, this.token.address, loanAmount, '0x');
+      await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: ZERO_ADDRESS, to: receiver.address, value: loanAmount });
+      await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: receiver.address, to: ZERO_ADDRESS, value: loanAmount });
+      await expectEvent.inTransaction(tx, this.token, 'Transfer', { from: receiver.address, to: this.token.address, value: flashFee });
+      await expectEvent.inTransaction(tx, receiver, 'BalanceOf', { token: this.token.address, account: receiver.address, value: receiverInitialBalance.add(loanAmount) });
+      await expectEvent.inTransaction(tx, 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(receiver.address)).to.be.bignumber.equal(receiverInitialBalance.sub(flashFee));
+      expect(await this.token.balanceOf(this.token.address)).to.be.bignumber.equal(flashFee);
+      expect(await this.token.allowance(receiver.address, this.token.address)).to.be.bignumber.equal('0');
+    });
+  });
+});