[SCH] Smart Contract Hacking 16편 - DeFi : Dex

0xDave·2023년 5월 6일
0

Ethereum

목록 보기
109/112
post-thumbnail

Task1


유니스왑V2 컨트랙트를 참고해서 빈칸들을 채워보자.

pragma solidity ^0.8.13;

import "@openzeppelin/contracts/utils/Address.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import "@openzeppelin/contracts/access/Ownable.sol";

import "../interfaces/IUniswapV2.sol";
interface IWETH is IERC20 {
    function deposit() external payable;
    function transfer(address to, uint value) external returns (bool);
    function withdraw(uint) external;
}

/**
 * @title Chocolate
 * @author JohnnyTime (https://smartcontractshacking.com)
 */
contract Chocolate is ERC20, Ownable {

    using Address for address payable;

    IUniswapV2Router02 public uniswapV2Router;

    address public weth;
    address public uniswapV2Pair;
    
    constructor(uint256 _initialMint) ERC20("Chocolate Token", "Choc") {

        // TODO: Mint tokens to owner
        
        // TODO: SET Uniswap Router Contract

        // TODO: Set WETH (get it from the router)

        // TODO: Create a uniswap Pair with WETH, and store it in the contract

    }

    /*
        @dev An admin function to add liquidity of chocolate with WETH 
        @dev payable, received Native ETH and converts it to WETH
        @dev lp tokens are sent to contract owner
    */
    function addChocolateLiquidity(uint256 _tokenAmount) external payable onlyOwner {

        // TODO: Transfer the tokens from the sender to the contract
        // Sender should approve the contract spending the chocolate tokens
        
        // TODO: Convert ETH to WETH

        // TODO: Approve the router to spend the tokens

        // TODO: Add the liquidity, using the router, send lp tokens to the contract owner
        
    }
    
    /*
        @dev An admin function to remove liquidity of chocolate with WETH 
        @dev received `_lpTokensToRemove`, removes the liquidity
        @dev and sends the tokens to the contract owner
    */
    function removeChocolateLiquidity(uint256 _lpTokensToRemove) external onlyOwner {
        
        // TODO: Transfer the lp tokens from the sender to the contract
        // Sender should approve token spending for the contract

        // TODO: Approve the router to spend the tokens
    
        // TODO: Remove the liquiduity using the router, send tokens to the owner
        
    }

    /*
        @dev User facing helper function to swap chocolate to WETH and ETH to chocolate
        @dev received `_lpTokensToRemove`, removes the liquidity
        @dev and sends the tokens to the contract user that swapped
    */
    function swapChocolates(address _tokenIn, uint256 _amountIn) public payable {

        // TODO: Implement a dynamic function to swap Chocolate to ETH or ETH to Chocolate
        
        if(_tokenIn == address(this)) {
            // TODO: Revert if the user sent ETH
            
            // TODO: Set the path array
            
            // TODO: Transfer the chocolate tokens from the sender to this contract

            // TODO: Approve the router to spend the chocolate tokens
            
        } else if(_tokenIn == weth) {
            // TODO: Make sure msg.value equals _amountIn

            // TODO: Convert ETH to WETH
            
            // TODO: Set the path array
            
            // TODO: Approve the router to spend the WETH
            
        } else {
            revert("wrong token");
        }

        // TODO: Execute the swap, send the tokens (chocolate / weth) directly to the user (msg.sender)
        
    }
}

1. constructor


    constructor(uint256 _initialMint) ERC20("Chocolate Token", "Choc") {

        // TODO: Mint tokens to owner
        
        // TODO: SET Uniswap Router Contract

        // TODO: Set WETH (get it from the router)

        // TODO: Create a uniswap Pair with WETH, and store it in the contract

    }

    constructor(uint256 _initialMint) ERC20("Chocolate Token", "Choc") {

        // TODO: Mint tokens to owner
        _mint(owner(), _initialMint);
        
        // TODO: SET Uniswap Router Contract
        uniswapV2Router = IUniswapV2Router02(0x7a250d5630B4cF539739dF2C5dAcb4c659F2488D);

        // TODO: Set WETH (get it from the router)
        weth = uniswapV2Router.WETH();

        // TODO: Create a uniswap Pair with WETH, and store it in the contract
        uniswapV2Pair = IUniswapV2Factory(uniswapV2Router.factory()).createPair(address(this), weth);
    }

문제에서 Router 주소를 제공해줘서 인터페이스를 적용해서 가져온 것 외에 특이한 점은 없다.


2. addChocolateLiquidity


    function addChocolateLiquidity(uint256 _tokenAmount) external payable onlyOwner {

        // TODO: Transfer the tokens from the sender to the contract
        // Sender should approve the contract spending the chocolate tokens
        
        // TODO: Convert ETH to WETH

        // TODO: Approve the router to spend the tokens

        // TODO: Add the liquidity, using the router, send lp tokens to the contract owner
        
    }

내가 썼던 답

    function addChocolateLiquidity(uint256 _tokenAmount) external payable onlyOwner {

        // TODO: Transfer the tokens from the sender to the contract
        IERC20(address(this)).transferFrom(msg.sender, uniswapV2Pair ,_tokenAmount);
        // Sender should approve the contract spending the chocolate tokens
        IERC20(address(this)).approve(address(this), _tokenAmount);

        // TODO: Convert ETH to WETH
        IWETH(weth).deposit{value: msg.value}();

        // TODO: Approve the router to spend the tokens
        IERC20(address(this)).approve(address(uniswapV2Router), _tokenAmount);

        // TODO: Add the liquidity, using the router, send lp tokens to the contract owner
        uniswapV2Router.addLiquidity(address(this), weth, _tokenAmount, msg.value, 1, 1, owner(), block.timestamp);

    }

모범답안

    function addChocolateLiquidity(uint256 _tokenAmount) external payable onlyOwner {
        // TODO: Transfer the tokens from the sender to the contract
        // Sender should approve the contract spending the chocolate tokens
        _transfer(msg.sender, address(this), _tokenAmount);

        // TODO: Convert ETH to WETH
        IWETH(weth).deposit{value: msg.value}();

        // TODO: Approve the router to spend the tokens
        IWETH(weth).approve(address(uniswapV2Router), msg.value);
        _approve(address(this), address(uniswapV2Router), _tokenAmount);

        // TODO: Add the liquidity, using the router, send lp tokens to the contract owner
        uniswapV2Router.addLiquidity(address(this), weth, _tokenAmount, msg.value, 1, 1, owner(), block.timestamp);
    }

우선 주석이 나눠져있는지 알았는데 하나였었고, ERC20을 상속했기 때문에 _transfer_approve를 바로 사용할 수 있었다. 한 가지 잘못 생각했던 점은 2 번째 라인에서 컨트랙트 스스로를 approve의 spender로 설정했다는 것이다. approve는 내가 갖고 있는 토큰을 다른 컨트랙트 혹은 다른 사용자에게 권한을 넘길 때 사용하는 함수이기 때문에 스스로를 spender로 설정할 일이 없다.


3. removeChocolateLiquidity


    function removeChocolateLiquidity(uint256 _lpTokensToRemove) external onlyOwner {
        
        // TODO: Transfer the lp tokens from the sender to the contract
        // Sender should approve token spending for the contract

        // TODO: Approve the router to spend the tokens
    
        // TODO: Remove the liquiduity using the router, send tokens to the owner
        
    }

내가 썼던 답

    function removeChocolateLiquidity(uint256 _lpTokensToRemove) external onlyOwner {
        
        // TODO: Transfer the lp tokens from the sender to the contract
        // Sender should approve token spending for the contract
        IUniswapV2Pair(uniswapV2Pair).transferFrom(msg.sender, address(this), _lpTokensToRemove);

        // TODO: Approve the router to spend the tokens
        IUniswapV2Pair(uniswapV2Pair).approve(address(uniswapV2Router), _lpTokensToRemove);
    
        // TODO: Remove the liquiduity using the router, send tokens to the owner
        uniswapV2Router.removeLiquidity(address(this), weth, _lpTokensToRemove, 1, 1, owner(), block.timestamp);
    }

모범답안

    function removeChocolateLiquidity(uint256 _lpTokensToRemove) external onlyOwner {
        // TODO: Transfer the lp tokens from the sender to the contract
        // Sender should approve token spending for the contract
        IERC20(uniswapV2Pair).transferFrom(owner(), address(this), _lpTokensToRemove);
        // TODO: Approve the router to spend the tokens
        IERC20(uniswapV2Pair).approve(address(uniswapV2Router), _lpTokensToRemove);
        // TODO: Remove the liquiduity using the router, send tokens to the owner
        uniswapV2Router.removeLiquidity(address(this), weth, _lpTokensToRemove, 1, 1, owner(), block.timestamp);
    }

IUniswapV2Pair 에도 transferFrom과 approve가 있는데 IERC20으로 해도 상관없을 것 같다. 이건 어떤 게 맞는지 정확히 모르겠다.


4. swapChocolates


    function swapChocolates(address _tokenIn, uint256 _amountIn) public payable {

        // TODO: Implement a dynamic function to swap Chocolate to ETH or ETH to Chocolate
        
        if(_tokenIn == address(this)) {
            // TODO: Revert if the user sent ETH
            
            // TODO: Set the path array
            
            // TODO: Transfer the chocolate tokens from the sender to this contract

            // TODO: Approve the router to spend the chocolate tokens
            
        } else if(_tokenIn == weth) {
            // TODO: Make sure msg.value equals _amountIn

            // TODO: Convert ETH to WETH
            
            // TODO: Set the path array
            
            // TODO: Approve the router to spend the WETH
            
        } else {
            revert("wrong token");
        }

        // TODO: Execute the swap, send the tokens (chocolate / weth) directly to the user (msg.sender)
        
    }

내가 썼던 답

   function swapChocolates(address _tokenIn, uint256 _amountIn) public payable {

        // TODO: Implement a dynamic function to swap Chocolate to ETH or ETH to Chocolate
        address[] memory path;
        
        if(_tokenIn == address(this)) {
            // TODO: Revert if the user sent ETH
            require(_tokenIn != address(weth), "It should be ETH");
            
            // TODO: Set the path array
            path[0] = _tokenIn;
            path[1] = weth;

            // TODO: Transfer the chocolate tokens from the sender to this contract
            IERC20(address(this)).transferFrom(msg.sender, address(this), _amountIn);

            // TODO: Approve the router to spend the chocolate tokens
            IERC20(address(this)).approve(address(uniswapV2Router), _amountIn);

            
        } else if(_tokenIn == weth) {
            // TODO: Make sure msg.value equals _amountIn
            require(msg.value == _amountIn);

            // TODO: Convert ETH to WETH
            IWETH(weth).deposit{value : msg.value}();
            
            // TODO: Set the path array
            path[0] = weth;
            path[1] = _tokenIn;
            
            // TODO: Approve the router to spend the WETH
            IERC20(weth).approve(address(uniswapV2Router), _amountIn);

            
        } else {
            revert("wrong token");
        }

        // TODO: Execute the swap, send the tokens (chocolate / weth) directly to the user (msg.sender)
        uniswapV2Router.swapExactTokensForTokens(_amountIn, 1, path, msg.sender, 0);
        
    }

두 번째 else if 부분에서 weth가 들어왔는데 왜 이더를 weth로 또 바꾸는지 이해가 안 갔었다. 그런데 곰곰이 생각해보니 이미 weth가 들어온 게 아니라 "weth를 chocolate 토큰으로 바꿀거예요"라고 요청하는 것이기 때문에 아직 weth가 안 들어왔다고 볼 수 있다.


모범답안

    function swapChocolates(address _tokenIn, uint256 _amountIn) public payable {
        // TODO: Implement a dynamic function to swap Chocolate to ETH or ETH to Chocolate
        address[] memory path = new address[](2);

        if (_tokenIn == address(this)) {
            // TODO: Revert if the user sent ETH
            require(msg.value == 0, "Don't send ETH!!!");
            // TODO: Set the path array
            path[0] = address(this);
            path[1] = weth;
            // TODO: Transfer the chocolate tokens from the sender to this contract
            _transfer(msg.sender, address(this), _amountIn);
            // TODO: Approve the router to spend the chocolate tokens
            _approve(address(this), address(uniswapV2Router), _amountIn);
        } else if (_tokenIn == weth) {
            // TODO: Convert ETH to WETH
            IWETH(weth).deposit{value: msg.value}();
            // TODO: Set the path array
            path[0] = weth;
            path[1] = address(this);
            // TODO: Approve the router to spend the WETH
            IWETH(weth).approve(address(uniswapV2Router), msg.value);
        } else {
            revert("wrong token");
        }

        // TODO: Execute the swap, send the tokens (chocolate / weth) directly to the user (msg.sender)
        uniswapV2Router.swapExactTokensForTokens(_amountIn, 0, path, msg.sender, block.timestamp);
    }
  1. 나는 path를 선언만 했었다. 이렇게 선언만 하고 initialize 하지 않은 array는 나중에 사용할 때 out-of-bounds error가 발생할 수 있다. 따라서 모범답안처럼 크기가 정해진 array를 initialize 해주는 게 좋다.

  2. 모범답안 마지막 부분에서 swapExactTokensForTokens의 두 번째 파라미터를 0으로 했는데, 최솟값을 0으로 한 이유는 조금 더 알아봐야겠다.


테스트 코드 짜기



내가 썼던 답

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

import "forge-std/Test.sol";
import "forge-std/console.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "src/dex-1/Chocolate.sol";
import {IUniswapV2Pair} from "src/interfaces/IUniswapV2.sol";

/**
@dev run "forge test --fork-url $ETH_RPC_URL --fork-block-number 15969633 --match-contract DEX1" 
*/

contract TestDEX1 is Test {

    address deployer;
    address user;
    Chocolate chocolate;
    address pair;
    address weth;

    uint256 initial_deployer_LP;


    function setUp() public {
        deployer = address(1);
        user = address(2);

        vm.deal(deployer, 100 ether);
        vm.deal(user, 100 ether);

        vm.prank(deployer);

        chocolate = new Chocolate(1000000);
        console.log("chocolate pair address :", address(chocolate));

        pair = chocolate.uniswapV2Pair();
        weth = chocolate.weth();

    }

    function test() public {
        vm.startPrank(deployer);
        chocolate.addChocolateLiquidity{value: 100 ether}(1000000);
        vm.stopPrank();

        console.log("deployer's LP token is :", IERC20(address(pair)).balanceOf(deployer));
        initial_deployer_LP = IERC20(address(pair)).balanceOf(deployer);

        assertEq(IERC20(address(chocolate)).balanceOf(user), 0);

        vm.startPrank(user);
        chocolate.swapChocolates{value: 10 ether}(weth, 10 ether);
        vm.stopPrank();

	    assertEq(IERC20(address(chocolate)).balanceOf(user), 100_000);

        vm.startPrank(user);
        chocolate.swapChocolates(address(chocolate), 100);
        vm.stopPrank();

        assertEq(IERC20(address(chocolate)).balanceOf(user), 100_000 - 100);

        vm.startPrank(deployer);
        uint256 halfLP = initial_deployer_LP / 2;
        IERC20(address(pair)).approve(address(chocolate), halfLP);
        chocolate.removeChocolateLiquidity(halfLP);
        vm.stopPrank();

        console.log("deployer's LP token is :", IERC20(address(pair)).balanceOf(deployer));

        assertEq(IERC20(address(chocolate)).balanceOf(deployer) >= 450000, true);

        assertEq(IERC20(address(weth)).balanceOf(deployer) >= 100 ether / 2, true);

    }

}

모범답안

contract TestDEX1 is Test {
    Chocolate chocolate;
    IUniswapV2Pair pair;

    address constant WETH_ADDRESS = address(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2);
    address constant RICH_SIGNER = address(0x8EB8a3b98659Cce290402893d0123abb75E3ab28);

    uint128 constant ETH_BALANCE = 300 ether;
    uint128 constant INITIAL_MINT = 1000000 ether;
    uint128 constant INITIAL_LIQUIDITY = 100000 ether;
    uint128 constant ETH_IN_LIQUIDITY = 100 ether;

    uint128 constant TEN_ETH = 10 ether;
    uint128 constant HUNDRED_CHOCOLATES = 100 ether;

    address deployer = makeAddr("deployer");
    address user = makeAddr("user");

    IERC20 weth = IERC20(WETH_ADDRESS);

    function setUp() public {
        vm.label(WETH_ADDRESS, "WETH");
        vm.label(RICH_SIGNER, "RICH_SIGNER");

        vm.deal(user, 100 ether);
        address richSigner = RICH_SIGNER;

        // Send ETH from rich signer to our deployer
        vm.prank(richSigner);
        (bool success, ) = deployer.call{value: ETH_BALANCE}("");
        require(success, "Transfer Failed!!!");
    }

    function test_Attack() public {
        /************************Deployment************************/
        // TODO: Deploy your smart contract to `chocolate`, mint 1,000,000 tokens to deployer
        vm.prank(deployer);
        chocolate = new Chocolate(INITIAL_MINT);
        // TODO: Print newly created pair address and store pair contract to `this.pair`
        address pairAddress = chocolate.uniswapV2Pair();
        //console.log(pairAddress);
        pair = IUniswapV2Pair(pairAddress);

        /************************Deployer add liquidity tests************************/
        // TODO: Add liquidity of 100,000 tokens and 100 ETH (1 token = 0.001 ETH)
        vm.startPrank(deployer);
        chocolate.approve(address(chocolate), INITIAL_LIQUIDITY);
        chocolate.addChocolateLiquidity{value: ETH_IN_LIQUIDITY}(INITIAL_LIQUIDITY);
        vm.stopPrank();
        // TODO: Print the amount of LP tokens that the deployer owns
        uint256 lpBalance = pair.balanceOf(deployer);
        console.log(lpBalance);

        /************************User swap tests************************/
        uint256 userChocolateBalance = chocolate.balanceOf(user);
        uint256 userWETHBalance = weth.balanceOf(user);

        // TODO: From user: Swap 10 ETH to Chocolate
        vm.prank(user);
        chocolate.swapChocolates{value: TEN_ETH}(address(weth), TEN_ETH);
        // TODO: Make sure user received the chocolates (greater amount than before)
        assertEq(chocolate.balanceOf(user) > userChocolateBalance, true);

        // TODO: From user: Swap 100 Chocolates to ETH
        vm.startPrank(user);
        chocolate.approve(address(chocolate), HUNDRED_CHOCOLATES);
        chocolate.swapChocolates(address(chocolate), HUNDRED_CHOCOLATES);
        vm.stopPrank();

        // TODO: Make sure user received the WETH (greater amount than before)
        assertEq(weth.balanceOf(user) > userWETHBalance, true);

        /************************Deployer remove liquidity tests************************/
        uint256 deployerChocolateBalance = chocolate.balanceOf(deployer);
        uint256 deployerWETHBalance = weth.balanceOf(deployer);

        // TODO: Remove 50% of deployer's liquidity
        vm.startPrank(deployer);
        pair.approve(address(chocolate), lpBalance / 2);
        chocolate.removeChocolateLiquidity(lpBalance / 2);
        vm.stopPrank();

        // TODO: Make sure deployer owns 50% of the LP tokens (leftovers)
        assertEq(pair.balanceOf(deployer), lpBalance / 2);

        // TODO: Make sure deployer got chocolate and weth back (greater amount than before)
        assertEq(weth.balanceOf(deployer) > deployerWETHBalance, true);
    }
}

피드백

  1. 먼저 contant로 이더 수를 정해놓고 시작하지 않아서 잘못된 숫자를 사용했다. 다음엔 꼭 테스트 코드 작성할 때 숫자를 먼저 표기하고 시작하자.
  2. vm.label()을 사용하면 테스트 결과에서 주소 대신 표시한 라벨로 나와서 보기 편하다.
  3. 모범답안에서는 deploy 단계부터 테스트를 시작했다. 다음에 테스트 코드 작성할 때는 deploy도 포함시키도록 하자. (+ 처음에 초콜릿을 민팅할 때 1,000,000을 넣었었는데 1000000 ether를 해줬어야 한다.)
  4. 모범답안의 코드가 보기 편하고 훨씬 깔끔하다. 다음엔 주석도 달아주자.
profile
Just BUIDL :)

1개의 댓글

comment-user-thumbnail
2023년 5월 9일

swapExactTokensForTokens의 두번째 파라미터 (amountOutMin)의 경우 슬리피지로 인해 해당 Min값 이하로 토콘이 out되게 되면 revert를 내게끔 되어있어 0을 넣어줍니다! ㅎㅎ

답글 달기