FlashLoan: Implementar ERC3156 en Exchange

contracts/Exchange.sol

@@ -3,15 +3,19 @@ pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/interfaces/IERC3156.sol";
 
 interface IRegistry {
     function getExchange(address _tokenAddress) external returns (address);
 }
 
-contract Exchange is ERC20 {
+contract Exchange is ERC20, IERC3156FlashLender {
     address public tokenAddress;
     address public registryAddress;
 
+    uint256 public lockedEthAmount;
+    uint256 public lockedTokenAmount;
+
     event TokenPurchase(address indexed buyer, uint256 indexed ethSold, uint256 indexed tokensBought);
     event EthPurchase(address indexed buyer, uint256 indexed tokensSold, uint256 indexed ethBought);
     event AddLiquidity(address indexed provider, uint256 indexed ethAmount, uint256 indexed tokenAmount);
@@ -31,11 +35,11 @@ contract Exchange is ERC20 {
         uint256 liquidity;
 
         if (getReserve() == 0) {
-            liquidity = address(this).balance;
+            liquidity = getBalance();
         } else {
             // Enforce the ratio once the pool is initialized to preserve
             // prices, but not before to allow initialization.
-            uint256 ethReserve = address(this).balance - msg.value;
+            uint256 ethReserve = getBalance() - msg.value;
             uint256 tokenReserve = getReserve();
             uint256 tokenAmount = (msg.value * tokenReserve) / ethReserve;
             require(_tokenAmount >= tokenAmount, "insufficient token amount");
@@ -58,7 +62,7 @@ contract Exchange is ERC20 {
         require(_amount > 0, "invalid amount");
 
         uint256 supply = totalSupply();
-        uint256 ethAmount = (address(this).balance * _amount) / supply;
+        uint256 ethAmount = (getBalance() * _amount) / supply;
         uint256 tokenAmount = (getReserve() * _amount) / supply;
 
         _burn(msg.sender, _amount);
@@ -71,7 +75,11 @@ contract Exchange is ERC20 {
     }
 
     function getReserve() public view returns (uint256) {
-        return IERC20(tokenAddress).balanceOf(address(this));
+        return IERC20(tokenAddress).balanceOf(address(this)) - lockedTokenAmount;
+    }
+
+    function getBalance() public view returns (uint256) {
+        return address(this).balance - lockedEthAmount;
     }
 
     function getAmount(
@@ -91,21 +99,24 @@ contract Exchange is ERC20 {
 
     function getTokenAmount(uint256 _ethSold) public view returns (uint256) {
         require(_ethSold > 0, "ethSold cannot be zero");
+        uint256 ethReserve = getBalance();
         uint256 tokenReserve = getReserve();
-        return getAmount(_ethSold, address(this).balance, tokenReserve);
+        return getAmount(_ethSold, ethReserve, tokenReserve);
     }
 
     function getEthAmount(uint256 _tokenSold) public view returns (uint256) {
         require(_tokenSold > 0, "tokenSold cannot be zero");
         uint256 tokenReserve = getReserve();
-        return getAmount(_tokenSold, tokenReserve, address(this).balance);
+        uint256 ethReserve = getBalance();
+        return getAmount(_tokenSold, tokenReserve, ethReserve);
     }
 
     function ethToToken(uint256 _minTokens, address recipient) private {
+        uint256 ethReserve = getBalance();
         uint256 tokenReserve = getReserve();
         uint256 tokensBought = getAmount(
             msg.value,
-            address(this).balance - msg.value,
+            ethReserve - msg.value,
             tokenReserve
         );
 
@@ -128,10 +139,11 @@ contract Exchange is ERC20 {
 
     function tokenToEthSwap(uint256 _tokensSold, uint256 _minEth) public {
         uint256 tokenReserve = getReserve();
+        uint256 ethReserve = getBalance();
         uint256 ethBought = getAmount(
             _tokensSold,
             tokenReserve,
-            address(this).balance
+            ethReserve
         );
 
         require(ethBought >= _minEth, "insufficient output amount");
@@ -161,10 +173,11 @@ contract Exchange is ERC20 {
         );
 
         uint256 tokenReserve = getReserve();
+        uint256 ethReserve = getBalance();
         uint256 ethBought = getAmount(
             _tokensSold,
             tokenReserve,
-            address(this).balance
+            ethReserve
         );
 
         IERC20(tokenAddress).transferFrom(
@@ -178,4 +191,72 @@ contract Exchange is ERC20 {
             msg.sender
         );
     }
+
+    function maxFlashLoan(address token) external view override returns (uint256) {
+        require(token == address(0) || token == tokenAddress, "ETH or Token lending only");
+        return token == address(0) ? getBalance() : getReserve();
+    }
+
+    function flashFee(address token, uint256 amount) external view override returns (uint256) {
+        require(token == address(0) || token == tokenAddress, "ETH or Token lending only");
+        require(amount > 0, "Borrow some first");
+        return 0;
+    }
+
+    function flashLoan(
+        IERC3156FlashBorrower receiver,
+        address token,
+        uint256 amount,
+        bytes calldata data
+    ) external override returns (bool) {
+        require(token == address(0) || token == tokenAddress, "ETH or Token lending only");
+
+        uint256 beforeEthReserve = getBalance();
+        uint256 beforeTokenReserve = getReserve();
+
+        uint256 beforeAmount = token == address(0) ? beforeEthReserve : beforeTokenReserve;
+
+        require(beforeAmount >= amount, "Not enough to borrow");
+
+        uint256 fee = 0;
+
+        uint256 lockTokenAmount = getAmount(
+            amount,
+            beforeEthReserve,
+            beforeTokenReserve
+        );
+        uint256 lockEthAmount = getAmount(
+            amount,
+            beforeTokenReserve,
+            beforeEthReserve
+        );
+
+        if (token == address(0)) {
+            lockedTokenAmount = lockedTokenAmount + lockTokenAmount;
+            payable(address(receiver)).transfer(amount);
+        } else {
+            lockedEthAmount = lockedEthAmount + lockEthAmount;
+            IERC20(tokenAddress).transfer(address(receiver), amount);
+        }
+
+        require(
+            receiver.onFlashLoan(msg.sender, token, amount, fee, data) == keccak256("ERC3156FlashBorrower.onFlashLoan"),
+            "IERC3156: Callback failed"
+        );
+
+        uint256 afterEthReserve = getBalance();
+        uint256 afterTokenReserve = getReserve();
+
+        uint256 afterAmount = token == address(0) ? afterEthReserve : afterTokenReserve;
+
+        require(afterAmount == beforeAmount + fee, "Give back the money");
+
+        if (token == address(0)) {
+            lockedTokenAmount = lockedTokenAmount - lockTokenAmount;
+        } else {
+            lockedEthAmount = lockedEthAmount - lockEthAmount;
+        }
+
+        return true;
+    }
 }

contracts/TestFlashBorrower.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/interfaces/IERC3156.sol";
+
+contract TestFlashBorrower is IERC3156FlashBorrower {
+    address public tokenAddress;
+
+    constructor(address _token) {
+        require(_token != address(0), "invalid token address");
+        tokenAddress = _token;
+    }
+
+    function onFlashLoan(
+        address initiator,
+        address token,
+        uint256 amount,
+        uint256 fee,
+        bytes calldata data
+    ) external override returns (bytes32) {
+        require(address(this).balance > fee, "Need ETH");
+        require(IERC20(tokenAddress).balanceOf(address(this)) > fee, "Need token");
+
+        if (token == address(0)) {
+            payable(msg.sender).transfer(amount + fee);
+        } else {
+            IERC20(tokenAddress).transfer(msg.sender, amount + fee);
+        }
+
+        return keccak256("ERC3156FlashBorrower.onFlashLoan");
+    }
+
+    receive() external payable {}
+}

test/exchange-test.js

@@ -110,4 +110,28 @@ describe("Exchange", function () {
     // // revisar estos valores
     // expect(await token.balanceOf(alice.address)).to.eq(expectedOutputForAlice);
   });
+
+  describe("flashLoan", function () {
+    let borrower;
+
+    beforeEach(async function () {
+      await token.approve(exchange.address, amountA);
+      tx = exchange.addLiquidity(amountA, { value: amountB });
+
+      const TestFlashBorrower = await ethers.getContractFactory("TestFlashBorrower");
+      borrower = await TestFlashBorrower.deploy(token.address);
+
+      await token.transfer(borrower.address, amountA);
+      await deployer.sendTransaction({to: borrower.address, value: amountB});
+    });
+
+    it("flash loan eth", async () => {
+      tx = await exchange.flashLoan(
+        borrower.address,
+        ethers.constants.AddressZero,
+        ethers.utils.parseEther("2"),
+        ethers.constants.HashZero
+      );
+    });
+  });
 });