Строим DEX: пул на Solidity и тесты, которые ловят атаку

Два урока подряд мы считали пул на Rust: инвариант и своп, потом диапазоны. Теперь тот же инвариант переезжает туда, где ему положено жить, в контракт. Урок проектный: к концу у тебя собирается пул, который принимает ликвидность, обменивает токены, выдаёт флеш-свопы и держит удар от живого атакующего. Проверять будем не глазами, а forge test, включая фаззинг инварианта.

Каркас проекта

Работаем в Foundry, который мы уже трогали в уроке про инструментарий и в уроке про безопасность.

forge init our-dex-contracts
cd our-dex-contracts
forge test

Раскладка простая: контракты в src/, тесты в test/, зависимости в lib/. Всё, что ниже, кладётся в два файла: src/OurPool.sol и test/OurPool.t.sol.

Интерфейсы и хранилище

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

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

/// @notice Получатель флеш-свопа: пул зовёт этот метод, отдав токены, но не
/// закончив своп. Вернуть долг обязан именно этот вызов.
interface IFlashCallback {
    function onFlashSwap(address caller, uint256 amount0, uint256 amount1, bytes calldata data)
        external;
}

contract OurPool {
    /// @notice Доли, сжигаемые навсегда при первом вкладе.
    uint256 public constant MINIMUM_LIQUIDITY = 1000;
    /// @notice Комиссия свопа в базисных пунктах.
    uint256 public constant FEE_BPS = 30;
    uint256 private constant BPS = 10_000;

    IERC20 public immutable token0;
    IERC20 public immutable token1;

    uint256 public reserve0;
    uint256 public reserve1;
    uint256 public totalShares;
    mapping(address => uint256) public sharesOf;

    /// @dev 1 это «свободно», 2 это «внутри вызова». Целые вместо bool, чтобы
    /// не платить за переход слота из нуля в единицу на каждом свопе.
    uint256 private _entered = 1;

    event Mint(address indexed to, uint256 amount0, uint256 amount1, uint256 shares);
    event Burn(address indexed from, uint256 amount0, uint256 amount1, uint256 shares);
    event Swap(address indexed sender, bool zeroForOne, uint256 amountIn, uint256 amountOut);
    event Flash(address indexed sender, uint256 amount0, uint256 amount1);

    error Reentrant();
    error ZeroAmount();
    error InsufficientLiquidity();
    error InsufficientShares();
    error SlippageExceeded(uint256 got, uint256 min);
    error KInvariantBroken();

    modifier lock() {
        if (_entered != 1) revert Reentrant();
        _entered = 2;
        _;
        _entered = 1;
    }

    constructor(IERC20 _token0, IERC20 _token1) {
        token0 = _token0;
        token1 = _token1;
    }

Четыре решения видны прямо здесь.

immutable для адресов токенов. Значение зашивается в код контракта при деплое, и чтение стоит как константа, а не как обращение к хранилищу. Пул привязан к паре навсегда, и это правильно: смена токена под ногами у поставщиков означала бы кражу.

Пользовательские ошибки вместо строк. error SlippageExceeded(uint256 got, uint256 min) дешевле строки в require и, что важнее, машиночитаема: тест ловит её по селектору, а не по тексту.

Замок целым числом, а не булевым. Изменение слота с нуля на ненулевое стоит заметно дороже, чем изменение ненулевого на другое ненулевое. Начав с единицы, мы платим меньше на каждом свопе. Приём взят из практики, и он один из немногих случаев, когда микрооптимизация оправдана.

Комиссия константой. Пул с фиксированной комиссией проще рассуждать: нет управляющего, который может её поднять.

Расчёт свопа

    /// @notice Выход свопа по формуле пула. Комиссия снимается со входа,
    /// деление округляется вниз: остаток остаётся в пуле.
    function getAmountOut(uint256 amountIn, uint256 reserveIn, uint256 reserveOut)
        public
        pure
        returns (uint256)
    {
        if (amountIn == 0) revert ZeroAmount();
        if (reserveIn == 0 || reserveOut == 0) revert InsufficientLiquidity();
        uint256 netIn = amountIn * (BPS - FEE_BPS);
        return (netIn * reserveOut) / (reserveIn * BPS + netIn);
    }

Та же формула, что на Rust в уроке 46, только без checked_mul: начиная с версии 0.8 Solidity проверяет переполнение сам и откатывает транзакцию. Функция pure, поэтому её можно звать снаружи бесплатно и считать котировку до отправки.

Внесение ликвидности

    function addLiquidity(uint256 amount0, uint256 amount1)
        external
        lock
        returns (uint256 minted)
    {
        if (amount0 == 0 || amount1 == 0) revert ZeroAmount();

        if (totalShares == 0) {
            uint256 root = _sqrt(amount0 * amount1);
            if (root <= MINIMUM_LIQUIDITY) revert InsufficientLiquidity();
            minted = root - MINIMUM_LIQUIDITY;
            totalShares = MINIMUM_LIQUIDITY;
        } else {
            uint256 by0 = (amount0 * totalShares) / reserve0;
            uint256 by1 = (amount1 * totalShares) / reserve1;
            minted = by0 < by1 ? by0 : by1;
        }
        if (minted == 0) revert ZeroAmount();

        // Эффекты до взаимодействий: состояние согласовано ещё до того, как
        // управление уйдёт в код токена.
        reserve0 += amount0;
        reserve1 += amount1;
        totalShares += minted;
        sharesOf[msg.sender] += minted;

        token0.transferFrom(msg.sender, address(this), amount0);
        token1.transferFrom(msg.sender, address(this), amount1);
        emit Mint(msg.sender, amount0, amount1, minted);
    }

Обрати внимание на порядок: сначала мы правим всё состояние, и только потом зовём внешний код. Это правило checks-effects-interactions, и оно спасает даже тогда, когда про замок забыли.

Токен это чужой контракт, а transferFrom это вызов чужого кода. Считать, что он просто передвинет числа, нельзя: он может позвать тебя обратно, и мы это сейчас проверим тестом.

Изъятие ликвидности и своп

    function removeLiquidity(uint256 shares)
        external
        lock
        returns (uint256 amount0, uint256 amount1)
    {
        if (shares == 0) revert ZeroAmount();
        if (sharesOf[msg.sender] < shares) revert InsufficientShares();

        amount0 = (shares * reserve0) / totalShares;
        amount1 = (shares * reserve1) / totalShares;
        if (amount0 == 0 || amount1 == 0) revert ZeroAmount();

        sharesOf[msg.sender] -= shares;
        totalShares -= shares;
        reserve0 -= amount0;
        reserve1 -= amount1;

        token0.transfer(msg.sender, amount0);
        token1.transfer(msg.sender, amount1);
        emit Burn(msg.sender, amount0, amount1, shares);
    }

    /// @notice Обменять один токен на другой с защитой по минимальному выходу.
    /// @param minAmountOut граница, ниже которой своп не нужен. Без неё
    /// транзакция беззащитна перед сэндвичем: её исполнят по любой цене.
    function swap(bool zeroForOne, uint256 amountIn, uint256 minAmountOut)
        external
        lock
        returns (uint256 amountOut)
    {
        (uint256 reserveIn, uint256 reserveOut) =
            zeroForOne ? (reserve0, reserve1) : (reserve1, reserve0);
        amountOut = getAmountOut(amountIn, reserveIn, reserveOut);
        if (amountOut < minAmountOut) revert SlippageExceeded(amountOut, minAmountOut);

        if (zeroForOne) {
            reserve0 += amountIn;
            reserve1 -= amountOut;
        } else {
            reserve1 += amountIn;
            reserve0 -= amountOut;
        }

        (IERC20 tokenIn, IERC20 tokenOut) = zeroForOne ? (token0, token1) : (token1, token0);
        tokenIn.transferFrom(msg.sender, address(this), amountIn);
        tokenOut.transfer(msg.sender, amountOut);
        emit Swap(msg.sender, zeroForOne, amountIn, amountOut);
    }

Параметр minAmountOut не удобство, а обязательная часть интерфейса. Пользователь считает котировку в момент отправки, а исполняется транзакция через несколько секунд и после чужих сделок. Без этой границы любой своп это подпись под фразой «согласен на любую цену», а желающие такой подписью воспользоваться описаны в уроке про MEV.