Graviton での RSA の高速化: 形式的検証で正当性を証明し開発も加速

本ブログは 2024 年 8 月 8 日に公開された Amazon Science Blog “Formal verification makes RSA faster — and faster to deploy” を翻訳したものです。

Amazon の Graviton2 チップ向け最適化で効率が向上し、形式的検証により開発時間も短縮しました。

オンラインにおける安全な取引のほとんどは、RSA のような公開鍵暗号方式によって保護されています。RSA のセキュリティは、大きな数の素因数分解が困難であることに基づいています。公開鍵暗号によって秘密鍵を安全に交換できるため、セキュリティが向上します。ただし、大きな整数のべき乗剰余演算などの重い計算処理が必要となるため、大きな計算オーバーヘッドが伴います。

研究者やエンジニアは公開鍵暗号の効率を高めるためにさまざまな最適化手法を導入してきましたが、最適化に伴う複雑さのため、暗号アルゴリズムの正しさを検証することが難しくなっています。暗号アルゴリズムのバグは致命的な結果をもたらしかねません。

この記事では、Amazon Automated Reasoning Group が Graviton2 チップにおける RSA 署名のスループットを鍵サイズに応じて 33～94% 向上させ、形式的検証を用いて最適化コードの機能的正当性を証明した取り組みについて紹介します。

Graviton2 は、Amazon Annapurna Labs が開発した Arm Neoverse N1 コアベースのサーバークラス CPU です。Graviton2 での RSA 署名のスループットを向上させるため、剰余演算を高速化するさまざまな手法と Graviton2 固有のアセンブリレベルの最適化を組み合わせました。最適化したコードの機能的正当性を証明するため、チームメンバーの John Harrison が開発した HOL Light 対話型定理証明器を用いて形式的検証を行いました。

コードは、秘密情報に依存する分岐やメモリアクセスパターンを排除する定数時間スタイルで記述されています。これは、関数の実行時間などの観測可能な情報から秘密情報を推測するサイドチャネル攻撃を防ぐためです。最適化した関数とその証明は、Amazon Web Services が提供する形式的検証済み多倍長整数演算ライブラリ s2n-bignum に含まれています。これらの関数は、AWS が管理する暗号ライブラリ AWS-LC と、そのバインディングである Amazon Corretto Crypto Provider (ACCP) および AWS Libcrypto for Rust (AWS-LC-RS) でも利用されています。

ステップ 1. Graviton2 での RSA の高速化

Graviton2 上で RSA アルゴリズムを最適化するには、乗算命令の配置と活用を慎重に検討する必要があります。64 ビット Arm CPU では、2 つの 64 ビット数値を乗算して最大 128 ビットの積を得る処理 (一般に 64×64→128 と表記) は、下位 64 ビットを返す MUL と上位 64 ビットを返す UMULH の 2 つの命令で構成されます。Graviton2 では、MUL のレイテンシは 4 サイクルで、発行後 2 サイクルの間乗算パイプラインをストールさせます。一方、UMULH のレイテンシは 5 サイクルで、発行後 3 サイクルの間乗算パイプラインをストールさせます。Neoverse N1 の乗算パイプラインは 1 つしかないのに対し加算パイプラインは 3 つあるため、乗算のスループットは 64 ビット加算の約 10 分の 1 にとどまります。

スループットを向上させるため、(1) 乗算命令を加算命令に置き換える別の乗算アルゴリズムを適用し、(2) SIMD (Single Instruction/Multiple Data) 命令を使用して乗算処理の一部を CPU のベクトルユニットにオフロードしました。

アルゴリズムの最適化

モンゴメリモジュラ乗算は、高速かつ安全な剰余演算を実現するために広く使われている手法です。モンゴメリ乗算は数値をモンゴメリ形式と呼ばれる特殊な形式で表現します。RSA アルゴリズムのように一連の剰余演算を実行する場合、中間積をモンゴメリ形式のまま保持することで、計算をより効率的に行えます。

モンゴメリ乗算を、多倍長整数乗算と独立したモンゴメリリダクションの組み合わせとして実装しました。これは 2 つの標準的な実装方法の 1 つです。

Graviton2 では、このアプローチを採用することで、よく知られたカラツバアルゴリズムを用いてコストの高い乗算を加算に置き換えることができます。カラツバアルゴリズムは、1 つの乗算を 3 つのより小さな乗算といくつかのレジスタシフトに分解する手法です。再帰的に適用でき、大きな数に対しては標準的な乗算アルゴリズムよりも効率的です。

カラツバアルゴリズムは、2,048 ビットや 4,096 ビットなどの 2 のべき乗のビットサイズに対して適用しました。それ以外のサイズ (例: 3,072 ビット) では、二次乗算を引き続き使用しています。また、カラツバ乗算は 2 つのオペランドが等しい場合にさらに最適化できるため、二乗演算に特化した関数も作成しました。

これらの最適化により、元のコードと比較して 2,048 ビットおよび 4,096 ビットの RSA 署名で 31～49% の高速化を達成しました。

マイクロアーキテクチャの最適化

多くの Arm CPU は Neon SIMD (Single Instruction/Multiple Data) アーキテクチャ拡張を実装しています。この拡張により、さまざまなサイズ (8/16/32/64 ビット) のベクトルとして扱える 128 ビットレジスタファイルと、それらのベクトルの一部または全部を並列処理できる SIMD 命令が追加されます。さらに、SIMD 命令はスカラ命令とは異なるパイプラインを使用するため、両者を並列に実行できます。

ベクトル化の戦略

ベクトル化とは、同じ演算の逐次的な繰り返しを、複数の値に対する単一の演算に置き換える手法で、一般的に処理効率が向上します。SIMD 命令を使用して、スカラの 64 ビット乗算をベクトル化しました。

多倍長整数乗算では、ベクトル化した 64 ビットの下位乗算コードが、スカラの 64 ビット上位乗算命令 (UMULH) とうまくオーバーラップしました。二乗演算では、2 つの 64×64→128 ビットの二乗演算をベクトル化することが効果的でした。モンゴメリリダクションにおける乗算では、64×64→128 ビットの乗算と 64×64→64 の下位乗算をベクトル化することが有効でした。ベクトル化するスカラ乗算を選定するにあたり、さまざまなベクトル化パターンを列挙して実行時間を計測するスクリプトを作成しました。短いコードフラグメントでは全パターンの列挙が可能でしたが、大きなコードフラグメントでは経験的な判断に頼る必要がありました。最終的な手法は、Seo et. al. at ICISC’14 に記述された手法をはじめ、他の代替手法との広範な比較実験を経て決定しました。

スカラユニットと SIMD ユニットは並列に動作しますが、整数レジスタと SIMD レジスタ間で入力値や中間結果を移動する必要が生じることがあり、これが大きな複雑さの要因となります。FMOV 命令を使えば 64 ビットスカラレジスタから SIMD レジスタにデータをコピーできますが、この命令はスカラ乗算器と同じパイプラインを使用するため、スカラ乗算のスループットが低下します。

代替手段として、まずベクトルレジスタにロードしてから MOV でスカラレジスタにコピーする方法もありますが、レイテンシは低いものの SIMD パイプラインを占有するため、SIMD 演算のスループットが低下します。直感に反しますが、最善の解決策は整数レジスタと SIMD レジスタに対して別々にメモリからロードし、それらの相対的な配置に注意を払うことでした。ただし、SIMD の結果が既に SIMD レジスタにある場合は、ストアロードによる往復よりも高速であるため、MOV 命令を使って一部の SIMD 結果を整数レジスタにコピーしました。

高速な定数時間テーブルルックアップコード

もう 1 つの独立した改善点として、高速なべき乗剰余アルゴリズム向けの定数時間ルックアップテーブルをベクトル化して再実装しました。この改善を先述の最適化と組み合わせることで、初期コードと比較して 2,048/4,096 ビット RSA 署名のスループットが 80～94% 向上し、3,072 ビット署名でも 33% の高速化を達成しました。

命令スケジューリング

Graviton2 はアウトオブオーダー CPU ですが、リオーダーバッファや発行キューなどのコンポーネント容量が有限であるため、パフォーマンスを最大化するには命令を慎重にスケジューリングすることが重要です。ここで述べた実装は手動の命令スケジューリングによるもので、良好な結果は得られましたが、かなりの時間を要しました。

制約求解と (簡略化された) マイクロアーキテクチャモデルに基づく SLOTHY スーパーオプティマイザを使用して、この作業を自動化する方法も検討しました。カラツバアルゴリズムで使用する一部の数値を事前計算するようにモンゴメリリダクションを追加調整し、SLOTHY による最適化を適用した結果、2,048/4,096 ビットのスループットで 95～120%、3,072 ビットで 46% の改善を実現しました。ただし、自動スケジューリングの正当性検証が困難であることが判明したため、この手法はまだ AWS-LC に組み込まれていません。スケジューリング最適化の正当性を自動的に証明する手法については、現在研究を進めています。

ステップ 2. コードの形式的検証

最適化したコードを本番環境にデプロイするには、正しく動作することを確認する必要があります。ランダムテストはシンプルな既知のケースを素早くチェックできる手軽なアプローチですが、より高い保証を得るために形式的検証を使用しています。このセクションでは、暗号プリミティブの機能的正当性を証明するために形式的検証をどのように適用するかを説明します。

s2n-bignum の紹介

AWS の s2n-bignum は、(1) x86-64 および Arm のアセンブリコードの形式的検証を行うためのフレームワークであり、同時に (2) そのフレームワークを使用して検証された暗号向けの高速なアセンブリ関数のコレクションです。

s2n-bignum における仕様

s2n-bignum のすべてのアセンブリ関数 (RSA で使用する新しいアセンブリ関数を含む) には、機能的正当性を記述した仕様が用意されています。仕様では、ある事前条件を満たすすべてのプログラム状態に対して、プログラムの出力状態がある事後条件を満たさなければならないと規定しています。例として、bignum_mul_4_8(uint64_t *z, uint64_t *x, uint64_t *y) は 2 つの 256 ビット (4 ワード) の数値を乗算して 512 ビット (8 ワード) の結果を生成する関数です。入力状態 s に対する事前条件 (省略版) は以下のとおりです。

  aligned_bytes_loaded s (word pc) bignum_mul_4_8_mc
∧ read PC s = word pc
∧ C_ARGUMENTS [z, x, y] s
∧ bignum_from_memory (x,4) s = a
∧ bignum_from_memory (y,4) s = b

これは、bignum_mul_4_8 のマシンコードがプログラムカウンタ PC の現在のアドレスにロードされていること (aligned_bytes_loaded)、C のアプリケーションバイナリインターフェイス (ABI) に従って関数引数にシンボリック値が割り当てられていること (C_ARGUMENTS …)、そしてシンボル a と b で表される多倍長整数がそれぞれ x と y の指すメモリ位置に 4 ワード分格納されていること (bignum_from_memory …) を意味します。

出力状態 s に対する事後条件 (省略版) は以下のとおりです。

bignum_from_memory (z,8) s = a * b

これは、乗算結果 a * b が位置 z から始まる 8 ワードのバッファに格納されることを意味します。

もう 1 つの構成要素として、入力状態と出力状態の間で満たすべき関係があります。

(MAYCHANGE_REGS_AND_FLAGS_PERMITTED_BY_ABI;
MAYCHANGE [memory :> bytes(z,8 * 8)]) (s_in,s_out)

これは、コードの実行によって ABI で許可されたレジスタやフラグ、および z から始まる 8 ワードのバッファは変更される可能性がありますが、それ以外の状態はすべて変更されないことを意味します。

HOL Light を用いたアセンブリの検証

実装が仕様に対して正しいことを証明するために、HOL Light 対話型定理証明器を使用しています。「ブラックボックス」型の自動定理証明器とは異なり、HOL Light のようなツールは、定型的な証明ステップの自動化と、ユーザーによる明示的かつプログラム可能なガイダンスのバランスを重視しています。紙の上や頭の中に証明があれば、熟練したユーザーはそれを対話型定理証明器で効率的に記述できます。s2n-bignum では、プログラムの検証に以下の 2 つの戦略を組み合わせています。

記号実行

記号実行では、特定の値の代わりにシンボリック変数を使用して入力プログラム状態を表現し、コード断片の実行後のシンボリック出力状態を推論します。これは実質的に、プログラム実行をより厳密かつ一般化した形で行うことに相当します。事後条件の証明は依然として必要ですが、プログラム実行に伴うアーティファクトが取り除かれ、純粋な数学的問題に帰着されます。

フロイド-ホーア論理スタイルの中間アノテーション

各中間アサーションは、先行するコードの事後条件であると同時に、後続するコードの事前条件として機能します。アサーションには、対応する事後条件を証明するために必要な詳細のみを含めれば十分です。この抽象化により、自動推論の処理能力の面でも、人間が結果を理解しやすくなるという面でも、記号シミュレーションがより扱いやすくなります。

Arm ハードウェアが s2n-bignum のモデルどおりに動作することを前提としていますが、モデルは慎重に開発されており、ハードウェアとの広範なクロスチェックによって検証されています。

形式的検証の今後の改善

s2n-bignum の形式的検証は、コードの実行時間などのサイドチャネルを通じた情報漏洩の有無といった、実装の非機能特性についてはまだカバーしていません。現在は、規律ある実装スタイルによって対処しています。具体的には、除算のように実行時間が変動する命令を使用せず、秘密データに依存する条件分岐やメモリアクセスパターンを排除しています。さらに、シンプルな静的チェックによってこれらの特性の一部を簡易検証し、ビット密度が大きく異なる入力でコードを実行して実行時間を分析し、予期しない相関がないか調査しています。

これらの規律と簡易チェックは AWS の標準的な手法であり、本記事で説明したすべての新しい実装に適用しています。情報漏洩がないことの形式的検証については、現在研究を進めています。