Fixing security vulnerabilities with AI

セキュリティの脆弱性をAIで修正する

Image of Ishikawa Setsuna

2023年11月、AIを活用してユーザーのコードベースのセキュリティ脆弱性の修正を提案する、コードスキャンautofixの開始を発表しました。この投稿では、autofixがどのように機能するのか、またテストと反復に使用する評価フレームワークについて説明します。

コードスキャンによる自動修正とは?

GitHub のコードスキャンは、リポジトリ内のコードを分析してセキュリティ脆弱性やその他のエラーを見つけます。スキャンは、スケジュールや、ブランチへのプッシュやプルリクエストのオープンなど、指定したイベントに応じてトリガーすることができます。問題が特定されると、ユーザーにアラートが表示されます。コードスキャンは、オープンソースやプライベートなツールを含め、ファーストパーティやサードパーティのアラートツールと一緒に使うことができます。GitHubは、当社のセマンティックコード解析エンジンであるCodeQLを利用したファーストパーティのアラートツールを提供しており、コードベースをデータのように照会することができます。私たちの社内のセキュリティ専門家は、一般的な言語やフレームワークのホストにわたってセキュリティの脆弱性を検出するためのクエリの豊富なセットを開発しました。この検出機能の上に構築されたコードスキャン自動修正機能は、AIが生成したアラートに対する修正を提案することで、セキュリティをさらにもう一段高めます。最初のイテレーションでは、JavaScriptとTypeScriptのアラートから始まり、プル・リクエストで検出されたCodeQLのアラートに対して自動修正が有効になります。問題とその修正策を自然言語で説明し、提案された修正をプルリクエストページに直接表示し、開発者が提案をコミット、却下、または編集できるようにします。

CodeQLのようなコード解析ツールが問題を検出すると、影響を受けるコードと問題の説明を大規模言語モデル(LLM)に送信し、コードの機能を変更することなく問題を修正するコード編集を提案するよう依頼します。以下のセクションでは、LLMプロンプトの構築、モデルのレスポンスの処理、機能の品質評価、ユーザーへの提供の詳細と微妙な点について掘り下げていきます。

自動修正プロンプト

私たちの技術の中核には、LLMプロンプトを通して表現されるLLMへの要求があります。CodeQL静的解析は脆弱性を検出し、問題のあるコードの場所とその他の関連する場所を参照するアラートを生成します。たとえば、SQLインジェクションの脆弱性の場合、アラートは汚染されたデータがデータベースクエリを構築するために使用される場所にフラグを立て、さらに信頼できないデータがサニタイズされずにこの場所に到達する可能性を示す1つ以上のフローパスを含みます。私たちは警告から情報を抽出し、以下からなるLLMプロンプトを作成します:

  • この種の脆弱性に関する一般的な情報(通常、脆弱性の一般的な例とその修正方法を含み、CodeQLクエリーヘルプから抽出されます)。
  • 警告メッセージのソースコードの場所と内容。
  • フロー パスに沿ったすべての場所からの関連するコード スニペットと、警告メッセージで参照されているコードの場所。
  • 期待するレスポンスの指定。

次に、脆弱性を修正するためにコードを編集する方法を示すよう、モデルに要求します。

自動処理を可能にするために、モデル出力のための厳密なフォーマットを記述します。モデルは以下のセクションからなるMarkdownを出力します:

  1. 脆弱性を修正するための自然言語による詳細な指示。
  2. プロンプトで定義されたフォーマットに従った、必要なコード編集の完全な仕様。
  3. 該当する場合、プロジェクトに追加されるべき依存関係のリスト。これは、例えば、プロジェクトがまだ依存していないサードパーティのサニタイズ・ライブラリを修正で使用する場合に必要です。

自然言語による説明は、コードのスキャン警告とともにユーザーに表示され、その後、コードの編集と追加された依存関係から構築された差分パッチが表示されます。ユーザは提案された修正を確認し、必要に応じて編集・調整し、プルリクエストのコミットとして適用することができます。

前処理と後処理

私たちのゴールが素晴らしいデモを作成することであれば、このシンプルなセットアップで十分でしょう。しかし、現実世界の複雑さをサポートし、LLMの制限を克服するには、注意深いプロンプトの作成と後処理のヒューリスティックを組み合わせる必要があります。私たちのアプローチをこの投稿内で全てを説明することはできませんが、以下に影響の大きい側面のいくつかを概説します。

モデルを表示するコードの選択

CodeQL のアラートには、アラートの位置情報と、ソースからシンクまでのデータフローパスに沿ったステップが含まれることがあります。アラートメッセージでは、追加のソースコードロケーションが参照されることもあります。これらの場所のいずれかが、脆弱性を修正するための編集を必要とする場合があります。テスト・スイートのようなコードベースのさらなる部分も編集を必要とするかもしれませんが、プロンプトの長さの制約のため、最も可能性の高い候補に焦点を当てます。

これらのコード位置のそれぞれについて、一連のヒューリスティックを使って、コード行数を最小にしつつ、必要なコンテキストを提供する周辺領域を選択し、目標の長さを達成するために必要に応じて関連性の低い部分を除外します。この領域は、ファイルの先頭にあるインポートや定義を含むように設計されており、これらは修正提案で補強する必要があることが多いためです。CodeQLアラートからの複数の場所が同じファイルに存在する場合、それらすべてに必要なコンテキストを与える結合コードスニペットを構成します。

その結果、1つ以上のコードスニペット(複数のソースコードファイルからの可能性もあります)のセットができ、編集が最も必要と思われるプロジェクトの部分をモデルに示し、モデルプロンプトとモデルレスポンスの両方で特定の行を参照できるように行番号が追加されます。捏造を防ぐために、プロンプトに含まれるコードに対してのみ編集を行うようにモデルを明示的に制約しています。

依存関係の追加

いくつかの修正では、データ・サニタイジング・ライブラリのような新しいプロジェクトの依存関係を追加する必要があります。そのためには、プロジェクトの依存関係を列挙した設定ファイルを見つけ、必要なパッケージがすでに含まれているかどうかを判断し、含まれていなければ必要な追加を行う必要があります。これらのすべてのステップにLLMを使うこともできますが、その場合、LLMにコードベース内のファイルのリストと関連するファイルの内容を見せる必要があります。これではモデルの呼び出し回数もプロンプトトークンの数も増えてしまいますので、その代わりに、修正に使われる外部の依存関係をリストアップするようにモデルに要求するだけです。私たちは言語固有のヒューリスティックを実装して、関連する設定ファイルを見つけ、それを解析して必要な依存関係がすでに存在するかどうかを判断し、存在しない場合は必要な編集を生成する差分パッチに追加します。

コード編集のフォーマットの指定

コード編集を指定するモデルのためのコンパクトなフォーマットが必要です。一番わかりやすいのは、モデルに標準的なdiffパッチを直接出力させることでしょう。残念なことに、実験によると、この方法はモデルの既知の算術的な難点を悪化させ、発見的な修正を行うのに十分なコードコンテキストがないまま、しばしば誤った行番号の計算を行うことがわかりました。私たちは、モデルが使用できる行編集コマンドの固定セットを定義するなど、いくつかの代替案を実験しました。実際に最良の結果をもたらしたアプローチは、モデルが「変更前」と「変更後」のコードブロックを提供し、変更が必要なスニペット(周囲のコンテキスト行を含む)と編集内容を示すというものです。

モデルエラーの克服

モデル出力の小さなエラーを検出して修正するために、さまざまな後処理ヒューリスティックを採用しています。例えば、”before “コードブロックは元のソースコードと完全に一致しないかもしれないし、行番号がわずかにずれているかもしれません。私たちは、インデント、セミコロン、コードコメントなどのエラーを克服し修正しながら、元のコードと一致するようにファジー検索を実装しています。編集したコードに構文エラーがないか、パーサーを使ってチェックします。また、名前解決チェックや型チェックなどのセマンティック・チェックも行います。ヒューリスティックに修正できないエラーを検出した場合は、提案された編集を(部分的に)間違っているとフラグを立てます。モデルがプロジェクトに追加する新しい依存関係を提案する場合、これらのパッケージがエコシステムのパッケージレジストリに存在することを検証し、既知のセキュリティ脆弱性や悪意のあるパッケージがないかチェックします。

評価と反復

LLMの計算コストを最小化すると同時に、プロンプトとヒューリスティックを繰り返し改善するために、修正提案を大規模に評価する必要があります。autofixをデモクオリティから本番クオリティに引き上げるにあたり、私たちは迅速な評価と反復を可能にする広範な自動テストハーネスに依存しました。

テストハーネスの最初のコンポーネントは、コードスキャンアラートを含むオープンソースリポジトリを処理するデータ収集パイプラインであり、アラート位置のテストカバレッジを持つアラートを収集する。最初にサポートされた言語であるJavaScript / TypeScriptでは、63のCodeQLクエリからテストカバレッジを持つ1,400以上のアラートを収集しました。

テストハーネスの2つ目のコンポーネントは、評価セット内の各アラートに対してautofixを実行するGitHub Actionsワークフローです。生成された修正をフォークでコミットした後、ワークフローはCodeQLとリポジトリのテストスイートの両方を実行し、修正の妥当性を評価します。特に、修正は以下の場合にのみ成功したとみなされます:

  • CodeQLアラートを削除する。
  • 新しい CodeQL 警告が発生しない。
  • 構文エラーが発生しない。
  • リポジトリテストの結果を変更しない。

プロンプト、コード編集形式、およびさまざまな後処理ヒューリスティックについて繰り返し検討しながら、このテスト ハーネスを使用して、変更が成功率を向上させていることを確認しました。自動化された評価と定期的な手作業によるトリアージを組み合わせることで、最も一般的な問題に労力を集中させ、自動化されたフレームワークの精度を検証しました。データ駆動型開発へのこの厳格なアプローチにより、成功率を3倍に高めると同時に、LLMの計算要件を6分の1に減らすことができました。

アーキテクチャ、インフラ、ユーザー・エクスペリエンス

有用な修正を生成することは第一歩ですが、それをユーザーに提供するには、フロントエンドとバックエンドのさらなる改良が必要です。シンプルさを追求し、可能な限り既存の機能の上にautofixを構築しました。ユーザーエクスペリエンスは、プルリクエストのコードスキャニングのエクスペリエンスを向上させます。コードスキャンアラートとともに、ユーザは修正案を見ることができます。修正案には、プルリクエストの差分の範囲外の複数のファイルの変更案が含まれることもあります。修正に関する自然言語による説明も表示されます。ユーザーは提案された修正をプルリクエストに直接コミットしたり、ローカルの IDE や GitHub Codespacesで編集したりできます。

バックエンドも、既存のコードスキャンインフラストラクチャの上に構築されているため、ユーザーにとってシームレスです。お客様は、サポートされているCodeQLクエリの修正候補を見るために、コードスキャンワークフローに変更を加える必要はありません。

Diagram outlining the code scanning pull request workflow.

ユーザーはプルリクエストを開くか、コミットをプッシュします。コードスキャンは通常通り、アクションワークフローまたはサードパーティ CI システムのワークフローの一部として実行され、結果を SARIF フォーマットでコードスキャン API にアップロードします。コードスキャンバックエンドサービスは、結果がサポートされている言語のものかどうかをチェックし、もしそうであれば、修正ジェネレーターをCLIツールとして実行します。修正ジェネレーターは、SARIFアラートデータを活用し、リポジトリから関連するソースコードを追加して、LLMのプロンプトを作成します。LLMは、Azure上でLLMを実行する内部展開されたAPIへの認証されたAPIコールによって呼び出されます。LLMのレスポンスは、特定のクラスの有害なレスポンスを防止するためのフィルタリングシステムにかけられます。修正ジェネレーターは、LLM レスポンスを後処理して修正案を生成します。コードスキャンバックエンドは結果のサジェスチョンを保存し、プルリクエストビューでアラートとともにレンダリングできるようにします。サジェスチョンは可能な限り再利用できるようにキャッシュされ、LLM の計算要件を削減します。

すべてのGitHub製品と同様に、私たちは標準的かつ内部的なセキュリティ手順に従い、ユーザーを保護するために厳格なセキュリティとプライバシーのレビュープロセスを経てアーキテクチャ設計を行いました。また、プロンプト・インジェクション攻撃などのAI特有のリスクに対しても予防策を講じました。ソフトウェアのセキュリティが完全に保証されることはありませんが、私たちはレッドチームテストを実施し、モデル・レスポンス・フィルターやその他の安全メカニズムをストレステストし、セキュリティ、有害なコンテンツ、モデルの偏りに関するリスクを評価しました。

遠隔測定とモニタリング

autofixをローンチする前に、私たちはパフォーマンスをモニターし、その影響を測定できるようにしたいと考えました。プロンプトやモデルのレスポンスには、プライベートなユーザーコードが含まれている可能性があるため、私たちはプロンプトやモデルのレスポンスを収集しません。その代わりに、修正提案が生成されたアラートの割合、そのままブランチにコミットされた提案の割合、GitHub CLI や Codespaces を通じて適用された提案の割合、却下された提案の割合、提案のあるアラートとないアラートの修正率など、提案された修正に関するユーザーとのインタラクションを匿名化して集計したテレメトリを収集しています。より多くのユーザーがベータプログラムに参加するにつれて、私たちは提案の有用性を理解するためにこの遠隔測定を見ていきます。

さらに、AzureモデルAPIのオーバーロードや有害なコンテンツをブロックするフィルタのトリガーなどのエラーについてサービスを監視しています。autofixを無制限のパブリックベータに拡大し、最終的には一般利用可能にする前に、一貫した安定したユーザーエクスペリエンスを確保したいと考えています。

次はどうなるのか?

コード・スキャンのautofixベータ版をより多くのユーザーに展開しながら、私たちはフィードバックを収集し、ペーパーカットを修正し、私たちの提案が実際に世の中に出回っているセキュリティ脆弱性に対して有用であることを確認するためにメトリクスを監視しています。並行して、autofixをより多くの言語とユースケースに拡大し、ユーザー体験を改善しています。パブリック・ベータに参加したい方は、こちらからサインアップしてください。近日中のアップデートにご期待ください!


CodeQLのパワーを活用してください。今すぐ始めましょう。


The postFixing security vulnerabilities with AIappeared first onThe GitHub Blog.