Wasm モジュールの関数を Lean で形式検証する

最近、生成AIとLeanでどんどん数学研究が進んでいますね。

そうした動きを見て、ミーハーなのでLeanとWebAssemblyを組み合わせられないか、思考停止で考えてみました。WebAssemblyなら、さまざまな言語で書いたプログラムを共通のランタイムで動かせます。そこで、Wasmに形式検証を適用できれば、異なる言語で実装した同じアルゴリズムにも共通の検証ロジックを使えるのではないかと夢を見たくなります。

といっても、WebAssemblyは仮想的な命令セットで構成されたバイナリです。
そのバイナリを形式検証できるのでしょうか。

そこでこの記事では、WATからLeanの関数定義を作り、バイナリから読み出した関数本体と照合したうえで、形式検証を試します。

形式検証とは

形式検証は、プログラムの振る舞いを数学的なモデルで表し、仕様に書いた性質が成り立つかを機械的に確かめる方法です。代表的な方法に、状態遷移を探索するモデル検査や、定理証明器で論理的な証明を構成する演繹的検証があります。この記事では、Leanを定理証明器として使う演繹的検証を試します。

大規模な実例として、Anthropic は Theorem と協力し、OSS-Fuzz の 123 ライブラリにある 239 の実行入口を対象にした検証を報告しています。単一スレッドの静的 RISC-V バイナリについて、Linux のモデル上で 64 MB 未満の入力ではクラッシュしないことを Claude に証明させ、その証明を Rocq カーネルで検査したといいます。

この記事でも、ソースコードではなくコンパイル後のバイナリに注目します。対象を WebAssembly に絞り、Lean の実行モデルとエスクローの業務仕様を結び付けて性質を確かめる流れを、小さな例で追います。

形式検証やLeanについて知るための資料も挙げておきます。

エスクローの状態遷移を検証する

例題は、買い手と売り手を仲介するオンラインマーケットプレイスの決済管理です。買い手が代金を払うと、プラットフォームは取引完了まで決済サービス上で代金を預かります。商品受領などが確認できたら売り手への支払いを、取引不成立なら買い手への返金を上位処理に指示します。同じ取引で両方を指示すれば二重精算になるため、決済指示は片方だけにする必要があります。

実際の資金移動は上位の決済処理が担い、このサンプル関数は状態を更新して結果コードを返します。同じ状態遷移を C と Rust で実装し、安全版と、支払い後も未決済状態を残す脆弱版を用意しました。

状態は 1 バイトで表し、state[0] に取引の段階を保存します。操作は入金、売り手への支払い、買い手への返金の 3 種類です。

安全版では、FUNDED のときだけ支払いまたは返金を受け付け、処理後は終端状態に移ります。脆弱版は売り手に支払った後も状態を FUNDED のままにするため、その後の返金も受け付けてしまいます。

この状態遷移をCとRustで実装したコードが次のとおりです。

#include <stdint.h>

enum {
    EMPTY = 0, FUNDED = 1, RELEASED = 2, REFUNDED = 3,
    FUND = 0, RELEASE = 1000, REFUND = 1000000
};

uint32_t escrow_step(uint8_t *state, uint32_t action) {
    uint32_t phase = state[0];

    if (phase == EMPTY && action == FUND) {
        state[0] = FUNDED;
        return 1;
    }
    if (phase == FUNDED && action == RELEASE) {
#ifndef VULNERABLE
        state[0] = RELEASED;
#endif
        return 2;
    }
    if (phase == FUNDED && action == REFUND) {
        state[0] = REFUNDED;
        return 3;
    }
    return 0;
}

#![no_std]

use core::panic::PanicInfo;

#[panic_handler]
fn panic(_: &PanicInfo) -> ! { loop {} }

#[no_mangle]
pub unsafe extern "C" fn escrow_step(
    state: *mut u8, action: u32,
) -> u32 {
    let phase = *state as u32;

    if phase == 0 && action == 0 {
        *state = 1;
        return 1;
    }
    if phase == 1 && action == 1000 {
        #[cfg(not(feature = "vulnerable"))]
        { *state = 2; }
        return 2;
    }
    if phase == 1 && action == 1000000 {
        *state = 3;
        return 3;
    }
    0
}

C は -DVULNERABLE、Rust は --cfg 'feature="vulnerable"' を付けると、支払い後に RELEASED を保存する処理が外れます。どちらも -O2 でコンパイルしました。

Wasm 関数に渡す操作コードは、入金が 0、売り手への支払いが 1000、買い手への返金が 1000000 です。

生成したモジュールは wasm-tools validate に通しました。サイズと SHA-256 は、今回 Lean に取り込んだ成果物の記録です。

形式検証する内容

以降では、仕様の定義から性質の確認までを「形式検証」、Lean が受理した個々の性質の証明を「証明」と呼びます。

操作列は FUND、RELEASE、REFUND からなる任意の有限列です。安全版の仕様は「その列を実行しても、売り手への支払いと買い手への返金を合わせて 2 回以上行わない」です。

脆弱版には具体的な反例があります。

結果 2 は売り手への支払い、3 は買い手への返金を表します。上位処理が両方の指示を実行すれば、同じ取引で二重精算になります。この3操作の反例は短く、ユニットテストでも検出できます。テストは試した操作列について結果を確認します。Leanではさらに、定義した3種類の操作からなる任意の有限列について、安全版が決済結果コードを最大1回しか返さないことを示します。

Wasm を共通の検証対象にする

C と Rust のように異なる言語で実装した関数をそれぞれ Wasmモジュールにし、共通の Lean 実行モデルとエスクロー仕様で検証します。ただし、対象にできるのは、実行モデルが対応する命令を使った関数です。ソースコード同士の同値性や、コンパイラがソースの意味を保存することまでは証明しません。

命令の実行を追うには、値スタックやメモリなどの状態を扱います。
それらがそのまま定理に現れると読みにくくなるのは自然です。
とはいえ、実行状態を追う検証は Wasm に限りません。CompCert は C とアセンブリの実行状態を対応づけ、seL4 は抽象仕様と C 実装を段階的につないでいます。定理の表面では業務上の性質を主張し、状態の細部は実行モデルの中に置けます。

検証対象を実際の .wasm に結び付けるため、関数本体を二つの経路で Lean に取り込みます。
一方は WAT から Lean の Func 定義を作り、もう一方はバイト列をデコーダで直接読みます。図は、デコードした関数本体と Func 定義の一致を確かめ、その関数を実行モデルとエスクロー仕様に結び付ける流れです。

wasm-tools print などで .wasm をテキスト形式の WAT にし、筆者が実装した Rust 製トランスパイラ wat2lean で関数本体を Lean の Func 定義に変換します。リポジトリには、変換器のほか、Wasm の構文、バイナリデコーダ、命令実行モデルを提供する Lean ライブラリも含まれます。

habakan/wat2leanWAT から Lean へのトランスパイラと WebAssembly の実行・デコードライブラリgithub.com

cargo run は wat2lean リポジトリのルートで実行します。

wasm-tools print rust-safe.wasm > rust-safe.wat
cargo run -- rust-safe.wat 0 rustEscrowSafeCode > rust-safe.lean

この抜粋は、安全版で売り手への支払いを指示する条件を示します。local 2 は現在の状態、local 1 は操作、local 3 は次の状態と結果コードです。まず状態が FUNDED（1）かを調べ、続いて操作が RELEASE（1000）かを調べます。両方が一致した場合だけ、2 を local 3 に設定します。

local.get 2       ;; local 2 (phase) をスタックに積む
i32.const 255     ;; マスク値 255 を積む
i32.and           ;; AND で phase を 0〜255 に絞る
i32.const 1       ;; FUNDED (= 1) を積む
i32.ne            ;; phase != 1 なら 1 を積む
br_if 1           ;; 条件が 1 なら深さ 1 の分岐先へ
block             ;; RELEASE の条件判定を囲む
  local.get 1     ;; local 1 (action) をスタックに積む
  i32.const 1000  ;; RELEASE (= 1000) を積む
  i32.ne          ;; action != 1000 なら 1 を積む
  br_if 0         ;; 条件が 1 ならこのブロックを抜ける
  i32.const 2     ;; 支払い結果と次状態の値を積む
  local.set 3     ;; 値を local 3 に入れる
  br 1            ;; 深さ 1 の分岐先へ進む
end               ;; ブロックを閉じる

次の部分では、local 3 を state の先頭 1 バイトへ書き戻し、同じ値を関数の返り値にします。

local.get 0     ;; state のポインタをスタックに積む
local.get 3     ;; 次状態をスタックに積む
i32.store8      ;; state の先頭 1 バイトに書き込む
end             ;; 制御ブロックを閉じる
local.get 3     ;; 結果を戻り値としてスタックに積む

wat2lean は WAT の命令を .localGet、.const、.bin、.store8 などからなる Func に変換します。一方、Lean の codeBody デコーダは元のバイト列から関数本体を読みます。各バイナリで関数 0 の型が (2, 1) であることと、デコードした本体が Func と一致することを確認します。

import WasmCode
open WasmCode WasmCode.Instr WasmShape

theorem rustEscrowSafeCodeMatchesBinary : isFunc (codeBody rustEscrowSafeBytes 0) rustEscrowSafeCode = true := by
  decide +kernel

theorem cEscrowSafeCodeMatchesBinary : isFunc (codeBody cEscrowSafeBytes 0) cEscrowSafeCode = true := by
  decide +kernel

theorem rustEscrowSafeTypeMatchesBinary : typeOf rustEscrowSafeBytes 0 = some (2, 1) := by
  decide +kernel

同じ型と関数本体の定理を、C と Rust の脆弱版にも置きました。WAT では import がなく、関数 0 が escrow_step として export されることを確認しています。ただし、Lean の codeBody は import がないモジュールを前提に関数 index を扱い、Lean の証明は import section と export section の対応までは確認していません。ハッシュ値は成果物を識別する記録で、Lean の証明対象ではありません。

Wasm 命令の実行規則を書く

命令を実行したとき、値やメモリがどう変わるかを定める規則をセマンティクスと呼びます。WebAssembly の仕様も実行状態と遷移規則を定義しています。たとえば i32.store8 は線形メモリに 1 バイトを書き、br は制御を指定したブロックへ移します。WebAssembly 仕様 1.0 の実行規則の一部に対応するモデルを Lean で定義し、命令列を実行します。

ここで使うライブラリは、命令の実行状態を Store と Out で表します。これは実行モデルの一部で、記事の定理が直接この型を並べる必要はありません。

namespace WasmCode

structure Store where
  mem : Nat -> UInt8
  size : Nat
  globals : List Nat

inductive Out where
  | normal (st ls : List Nat) (s : Store)
  | br (l : Nat) (st ls : List Nat) (s : Store)
  | ret (st ls : List Nat) (s : Store)
  | trap

end WasmCode

Store は線形メモリ、メモリサイズ、global 変数を持ちます。Out.normal は通常終了した実行結果で、値スタック、ローカル変数、実行後の Store が入ります。.trap はトラップ終了です。execInstr は命令を実行する関数で、引数の燃料が尽きた場合も .trap を返します。燃料は Lean の再帰実行に上限を付ける値で、Wasm モジュールが持つ設定ではありません。

具体的には、execInstr の型は (fs : Nat -> Option Fn) -> Nat -> Instr -> List Nat -> List Nat -> Store -> Out です。関数表、燃料、命令、値スタック、ローカル変数、Store を受け取り、実行結果を返します。

ただし、ライブラリが扱うのは WebAssembly 全体ではなく、明示した命令の部分集合です。対応していない命令はデコードに失敗するか、実行モデルで trap になります。定理はこのモデルに対するもので、モデルが仕様の全規則と一致することまでは主張しません。

実行前後の状態を仕様につなぐ

前節までで、Lean の Func 定義をバイナリ内の関数本体に結び付け、命令の実行モデルを見ました。ここでは Wasm 関数の一回の実行を、エスクロー仕様が扱う状態と結果コードに変換します。Lean の wasmEscrowStep は状態と操作を受け取り、execInstr で関数を呼び出し、通常終了した場合は実行後の状態と返り値を取り出します。

初期 Store ではエスクローの状態バイトをメモリ番地 2048 に置き、それ以外のメモリを 0 にします。C 向けは 258 ページと stack pointer の global [66576]、Rust 向けは 16 ページと global なしで初期化します。

ただし、.normal 以外の実行結果を wasmEscrowStep は (phase, 0) に変換します。そのため、一操作の一致定理が扱うのは、この関数が取り出した状態と結果コードです。仕様上「何もしない」入力で Wasm 関数がトラップしても、この一致だけでは検出できず、全呼び出しが正常終了することは証明していません。決済回数の定理も、wasmEscrowStep が返す結果コードを数えています。

実行モデルでは値スタックの先頭をリストの先頭に置くため、呼び出し時のスタックを [action, 2048] とします。関数呼び出しの処理が順序を戻し、Wasm 関数の local 0 に state ポインタ、local 1 に action を渡します。

仕様側の一操作は、たとえば次のように定義します。

def escrowSpecStep : EscrowPhase -> EscrowAction -> EscrowPhase × Nat
  | .empty, .fund => (.funded, 1)
  | .funded, .release => (.released, 2)
  | .funded, .refund => (.refunded, 3)
  | phase, _ => (phase, 0)

返り値 1 は入金、2 は売り手への支払い、3 は買い手への返金、0 は何もしなかったことを表します。脆弱版の仕様は支払い後も .funded のままにするため、FUND, RELEASE, REFUND の列で結果 2 と 3 の両方を返します。

この一操作の対応を任意の操作列へ帰納法で広げます。図は、Wasm の実行結果を仕様と照合し、安全版の決済結果コードが最大 1 回であることを示す流れです。

wasmEscrowStep が返す一操作の結果と仕様の遷移が一致することを、C と Rust の安全版・脆弱版でそれぞれ証明しました。その対応を任意長の操作列に帰納的に広げると、安全版について次の定理を得ます。

theorem cEscrowSafeTraceEmitsAtMostOneSettlementResult
    (actions : List EscrowAction) (phase : EscrowPhase) :
    settlementCount (cEscrowSafeTrace actions phase).1 <= 1 := by
  rw [cEscrowSafeTraceMatchesSpec]
  exact escrowSpecTraceSettlesAtMostOnce actions phase

theorem rustEscrowSafeTraceEmitsAtMostOneSettlementResult
    (actions : List EscrowAction) (phase : EscrowPhase) :
    settlementCount (rustEscrowSafeTrace actions phase).1 <= 1 := by
  rw [rustEscrowSafeTraceMatchesSpec]
  exact escrowSpecTraceSettlesAtMostOnce actions phase

ここで actions は FUND、RELEASE、REFUND からなる任意の有限列です。settlementCount は結果 2 と 3 を数えます。C と Rust の関数が返す決済指示の結果コードは、合わせて最大 1 回です。

脆弱版については、同じ Lean の実行モデルで具体的な反例を確認しました。

def attackTrace : List EscrowAction := [.fund, .release, .refund]

theorem cEscrowVulnerableTraceEmitsTwoSettlementResults :
    settlementCount (cEscrowVulnerableTrace attackTrace .empty).1 = 2 := by
  rw [cEscrowVulnerableTraceMatchesSpec]
  decide +kernel

theorem rustEscrowVulnerableTraceEmitsTwoSettlementResults :
    settlementCount (rustEscrowVulnerableTrace attackTrace .empty).1 = 2 := by
  rw [rustEscrowVulnerableTraceMatchesSpec]
  decide +kernel

最後に証明の前提を確認すると、EscrowKernel.lean の #print axioms では、安全版の定理が Lean の標準公理 propext と Quot.sound に依存すると表示されます。脆弱版の具体的な反例は propext のみです。sorry による未証明箇所は残っていません。

まとめ

• C と Rust の安全版・脆弱版を Wasm にコンパイルし、4 つの関数本体をバイナリデコーダの結果と照合しました。
• Lean の実行モデルと wasmEscrowStep が取り出した結果を使い、安全版が任意の有限操作列で二重の決済結果コードを返さないことを形式検証しました。
• 脆弱版では FUND, RELEASE, REFUND が二重の決済結果コードを返す反例を示しました。
• 対象は初期状態を指定した単一取引の 1 バイト状態、3 種類の操作、ライブラリが対応する命令に限られます。実際の資金移動・金額計算・並行処理は扱わず、.normal 以外を「何もしない」にまとめるため正常終了も保証しません。Wasm 全仕様への適合、コンパイラの正しさ、仕様自体が業務要件を満たすことも証明の対象外です。

この例では、コンパイル後の Wasm 関数を読みやすい状態遷移の仕様に結び付け、Lean で形式検証する流れを確認しました。