Add customizable fee receiver to ERC20FlashMint

contracts/mocks/ERC20FlashMintMock.sol

@@ -5,6 +5,9 @@ pragma solidity ^0.8.0;
 import "../token/ERC20/extensions/ERC20FlashMint.sol";
 
 contract ERC20FlashMintMock is ERC20FlashMint {
+    uint256 _flashFeeAmount;
+    address _flashFeeReceiverAddress;
+
     constructor(
         string memory name,
         string memory symbol,
@@ -13,4 +16,29 @@ contract ERC20FlashMintMock is ERC20FlashMint {
     ) ERC20(name, symbol) {
         _mint(initialAccount, initialBalance);
     }
+
+    function mint(address account, uint256 amount) public {
+        _mint(account, amount);
+    }
+
+    function setFlashFee(uint256 amount) public {
+        _flashFeeAmount = amount;
+    }
+
+    function flashFee(address token, uint256 amount) public view virtual override returns (uint256) {
+        super.flashFee(token, amount);
+        return _flashFeeAmount;
+    }
+
+    function setFlashFeeReceiver(address receiver) public {
+        _flashFeeReceiverAddress = receiver;
+    }
+
+    function flashFeeReceiver() public view returns (address) {
+        return _flashFeeReceiver();
+    }
+
+    function _flashFeeReceiver() internal view override returns (address) {
+        return _flashFeeReceiverAddress;
+    }
 }

contracts/token/ERC20/extensions/ERC20FlashMint.sol

@@ -43,6 +43,16 @@ abstract contract ERC20FlashMint is ERC20, IERC3156FlashLender {
         return 0;
     }
 
+    /**
+     * @dev Returns the receiver 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 receiver.
+     * @return The address for which the flash fee will be sent to.
+     */
+    function _flashFeeReceiver() internal view virtual returns (address) {
+        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 +83,14 @@ abstract contract ERC20FlashMint is ERC20, IERC3156FlashLender {
             receiver.onFlashLoan(msg.sender, token, amount, fee, data) == _RETURN_VALUE,
             "ERC20FlashMint: invalid return value"
         );
+        address flashFeeReceiver = _flashFeeReceiver();
         _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('flashFeeReceiver', function () {
+    it('default receiver', async function () {
+      expect(await this.token.flashFeeReceiver()).to.be.eq(ZERO_ADDRESS);
+    });
+  });
+
   describe('flashLoan', function () {
     it('success', async function () {
       const receiver = await ERC3156FlashBorrowerMock.new(true, true);
@@ -86,5 +92,53 @@ contract('ERC20FlashMint', function (accounts) {
       // _mint overflow reverts using a panic code. No reason string.
       await expectRevert.unspecified(this.token.flashLoan(receiver.address, this.token.address, MAX_UINT256, data));
     });
+
+    describe('custom flash fee & receiver', async function () {
+      const receiver = await ERC3156FlashBorrowerMock.new(true, true);
+      const receiverInitialBalance = new BN(200000);
+      const flashFee = new BN(5000);
+
+      before('init reciever balance & set flash fee',async function () {
+        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(receiver.address)).to.be.bignumber.equal(receiverInitialBalance);
+
+        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(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.add (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).sub(flashFee));
+        expect(await this.token.balanceOf(receiver.address)).to.be.bignumber.equal(receiverInitialBalance.sub(flashFee));
+        expect(await this.token.balanceOf(await this.token.flashFeeReceiver())).to.be.bignumber.equal('0');
+        expect(await this.token.allowance(receiver.address, this.token.address)).to.be.bignumber.equal('0');
+      });
+
+      it('custom flash fee receiver', async function () {
+        const flashFeeReceiverAddress = this.token.address;
+        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(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: flashFeeReceiverAddress, 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(flashFeeReceiverAddress)).to.be.bignumber.equal(flashFee);
+        expect(await this.token.allowance(receiver.address, flashFeeReceiverAddress)).to.be.bignumber.equal('0');
+      });
+    });
   });
 });