Как мы обучали LLM для поиска уязвимостей в смарт-контрактах Solidity22.11.2024 13:16

Наша команда в Positive Technologies занимается анализом безопасности смарт-контрактов, исследованием уязвимостей и разработкой инструментов для их обнаружения. Идея использовать LLM для анализа смарт-контрактов Solidity показалась крайне заманчивой. Загрузить код, запустить модель — и она сама находит уязвимости, генерирует отчет, а иногда даже предлагает исправления. Звучит отлично! Но, как показал мой опыт, между «звучит» и «работает» лежит огромная пропасть.

ChatGPT-4

Для начала мы решили изучить возможности существующих языковых моделей. Конечно, первым кандидатом стал ChatGPT. На момент начала экспериментов как раз вышла версия ChatGPT-4, и мы нашли интересный пример, как другая компания использовала ChatGPT для поиска багов в смарт-контрактах. Модель находила часть уязвимостей в простых контрактах и помогала с упрощением и пониманием кода в более сложных. Но было одно большое но: ChatGPT подходит только для публичных данных. Если код приватный, обратиться к ChatGPT для аудита не получится.

Поиск инструментов для анализа Solidity

Следующим шагом мы решили изучить, какие еще инструменты и плагины доступны для анализа смарт-контрактов на Solidity. Большинство из того, что нашли, было устаревшим и поддерживало старые версии Solidity, вроде 0.4, в которых многие баги еще не были исправлены на уровне компилятора (сейчас актуальна версия 0.8.29). Например, GPTLens тоже работает только с такими старыми версиями, что вызывает большое количество ложных срабатываний.

Несмотря на это, мы нашли пару инструментов, которые заслуживают внимания. Например, плагин для VS Code Solidity AI от Consensys Diligence, который может находить простые уязвимости и давать пояснения по коду. Этот инструмент полезен для анализа кода и обнаружения базовых ошибок, но на полноценный аудит не тянет. Мы также попробовали поиск багов через функцию Refactor Selected Code Blockв GitHub Copilot и его бесплатном аналоге Codeium, но их основная роль — помощь в написании кода, а не поиск уязвимостей. Для написания PoC под Solidity показали себя отлично.

Переход к созданию собственного LLM-агента

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

Основная проблема — датасет. Большинство доступных баз данных с уязвимостями устарели или содержат сплошные дубликаты. Например, многие из них включают старые контракты. Для обучения мы сверились с бенчмарками LLM-моделей и выбрали только вышедшую на тот момент Llama 3.1 — мощный вариант, но и требующий немало ресурсов.

Для расчета минимальных требований мы использовали калькулятор ресурсов LLM.

Затраты в GB видеопамяти для llama-3–8b

Затраты видеопамяти в ГБ для Llama-3–8B:

А вот затраты видео памяти на токены (рассчитаны по формулам из этой статьи).

В сумме получилось, что единственный приемлемый вариант, который мы могли позволить себе для обучения на ноутбуке, — это кванитизированная (ужатая) до 4bit llama-3–8b-bnb-4bit c 8k токенов, с суммарными затратами ≈6 ГБ видеопамяти.

Сначала обучали модель, используя фреймворк Unsloth на ноутбуке с 6 ГБ видеопамяти, но затем перешли на Google Colab с видеокартой Tesla 16 ГБ, что значительно ускорило процесс. У Google Colab есть свои недостатки: лимиты по времени на бесплатных сессиях, необходимость быть онлайн и периодически проходить капчу.

Промпт

Для дообучения методом LoRA используются датасеты с записями вида:

### Instruction:

{}

### Input:

{}

### Output:

{}

В этой задаче инструкция у нас была одна и та же, поэтому промпт выглядел так:

Below is a Solidity code snippet. Write a vulnerability analysis report for it.

### Code:

{}

### Response:

{}

Первые тесты и проблемы: «галлюцинации» и ложные срабатывания

Для качественного обучения модели нам нужен был хороший датасет. Мы изучили множество репозиториев уязвимостей для Solidity, но большинство данных касались старых версий. К счастью, ребята из Consensys предоставили доступ к своему датасету с описаниями багов и вариантами их исправления. После фильтрации и препроцессинга мы дообучили модель и получили первую версию. В тестовой выборке, где было 100 случайных контрактов, дообученная модель показала себя лучше, чем необученные версии, но всегда находила баги там, где их не было. Это еще раз показало необходимость доработки модели.

Чтобы уменьшить количество ложных срабатываний, мы добавили контракты 60 крупнейших по TVL проектов, таких как OpenZeppelin и Uniswap, как условно безопасные. Это заняло много времени, но позволило улучшить качество модели.

Во второй версии модели количество ответов false-positive сильно уменьшилось, но количество ответов false-negative увеличилось. Тогда мы вернули как было, поскольку лучше получить информацию о баге и отсеять ее на этапе проверки результатов, чем получить ответ, что «все хорошо», там, где баг есть.

Второй проблемой этой модели было зацикливание ответа:
### Response: function totalSupply() public view override returns (uint256) { return supply; } function supply() public view returns (uint256) { return totalSupply; } function totalSupply() public view override returns (uint256) { return supply; } function supply() public view returns (uint256) { return totalSupply; } function totalSupply() public view override returns (uint256) { return supply; }.....

Еще одна проблема — склонность модели к «галлюцинациям», то есть к генерации информации, которая не относится к запросу. Например, на один из запросов мы получили ответ, никак не связанный с аудитом:
### Response: Holmes is a detective who is known for his brilliant deductive reasoning and logical analysis. He is often portrayed as a master of disguise, using his skills to solve crimes and bring criminals to justice. Holmes is also known for his keen observation and ability to notice even the smallest details, which he uses to piece together clues and solve cases. He is often depicted as a loner, but he has a close relationship with his trusted friend and colleague, Dr. Watson. Holmes is a complex and fascinating character, and his story has been adapted for film, television, and literature for over a century.

Чтобы уменьшить количество ложных срабатываний, мы увеличили датасет и еще раз дообучили и протестировали модель. Помогло, но полностью проблему не решило. Возможно, свою роль сыграла квантизация (уменьшение точности параметров модели) до 4bit, что очень мало. Пришлось получать ответы по одному и тому же коду несколько раз и фильтровать.

Результаты

Ниже мы приводим результаты тестирования на 100 фрагментах контрактов тестовой выборки — не более трех ответов, и если хотя бы один из них давал правдивый результат, тесту добавлялся 1 балл.

Тестировались Llama 3–8B_4Q_K_M, Llama 3.1–8B-4Q_K_M, codellama-7B-4Q_K_M, ChatGPT-4o.

Для публичного кода можно использовать ChatGPT-4o, но в случае приватного аудита контрактов требуется использовать собственный агент. Правда, создание такого агента требует большого количества ресурсов и времени.

Если нужен надежный инструмент для аудита, придется уделить внимание валидации данных и созданию качественного датасета. Пока разработка собственного LLM-агента остается экспериментом, который, при грамотной настройке, может открыть новые возможности для автоматизации и повышения безопасности в блокчейне.

© Habrahabr.ru