[Ethernaut] Dex

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "openzeppelin-contracts-08/token/ERC20/IERC20.sol";
import "openzeppelin-contracts-08/token/ERC20/ERC20.sol";
import 'openzeppelin-contracts-08/access/Ownable.sol';

contract Dex is Ownable {
  address public token1;
  address public token2;
  constructor() {}

  function setTokens(address _token1, address _token2) public onlyOwner {
    token1 = _token1;
    token2 = _token2;
  }
  
  function addLiquidity(address token_address, uint amount) public onlyOwner {
    IERC20(token_address).transferFrom(msg.sender, address(this), amount);
  }
  
  function swap(address from, address to, uint amount) public {
    require((from == token1 && to == token2) || (from == token2 && to == token1), "Invalid tokens");
    require(IERC20(from).balanceOf(msg.sender) >= amount, "Not enough to swap");
    uint swapAmount = getSwapPrice(from, to, amount);
    IERC20(from).transferFrom(msg.sender, address(this), amount);
    IERC20(to).approve(address(this), swapAmount);
    IERC20(to).transferFrom(address(this), msg.sender, swapAmount);
  }

  function getSwapPrice(address from, address to, uint amount) public view returns(uint){
    return((amount * IERC20(to).balanceOf(address(this)))/IERC20(from).balanceOf(address(this)));
  }

  function approve(address spender, uint amount) public {
    SwappableToken(token1).approve(msg.sender, spender, amount);
    SwappableToken(token2).approve(msg.sender, spender, amount);
  }

  function balanceOf(address token, address account) public view returns (uint){
    return IERC20(token).balanceOf(account);
  }
}

contract SwappableToken is ERC20 {
  address private _dex;
  constructor(address dexInstance, string memory name, string memory symbol, uint256 initialSupply) ERC20(name, symbol) {
        _mint(msg.sender, initialSupply);
        _dex = dexInstance;
  }

  function approve(address owner, address spender, uint256 amount) public {
    require(owner != _dex, "InvalidApprover");
    super._approve(owner, spender, amount);
  }
}

우리의 mission은 contract에 있는 token1이나 token2의 token을 빼내는 것이다.

코드를 분석해 보면 간단한 DEX를 구현해 놓았다.

두 개의 token을 만들고, 유동성을 추가했다. 컨트랙트의 token1, token2의 양은 각각 100개이고 내가 가지고 있는 token1, token2는 각각 10개이다.

그리고 swap을 할 때 swapAmount 변수가 getSwapPrice() 함수를 호출한다.

getSwapPrice() 함수는 amount * to.balanceOf(address(this)/from.balanceOf(address(this))이다.

바꿀 토큰 * 컨트랙트에 있는 to의 토큰 / 컨트랙트에 있는 from의 토큰으로 구한다.

이 부분에서 취약점이 발생하는 데 solidity는 부동 소수점 처리를 안 한다. 그래서 만약 12.456라는 수가 있으면 소수점을 버림 처리하여 12가 된다. swap을 계속하면 getSwapPrice() 함수의 계산 값이 계속 바뀔 것이다.

컨트랙트와 내가 가지고 있는 token1, token2의 양은 그림과 같다.

그럼 swap(tk1, tk2, 10)을 한다면 10 * 100/100으로 10개가 swap 될 것이다.

token1 10개를 token2로 바꿨더니 DEX에 있는 token1은 110개, token2는 90개가 되었다.

그런데 어떻게 token2 20개를 token1로 바꾸면 24개가 될까?? getSwapPrice() 함수의 공식을 적용하면

20 * 110 / 90 = 24.444... 가 나온다. solidity는 부동 소수점 연산을 안 하니까 24로 내가 보유한 토큰이 20 ->24개로 증가하였다.

이런 식으로 계속 반복하다 보면 token이 증가하고 최종적으로 DEX에 있는 token1을 가져올 수 있다.