iddqd、あるいは最も難しい種類のunsafe Rust

 (oxide.computer)
1 ポイント 投稿者 GN⁺ 9 시간 전 | 1件のコメント | WhatsAppで共有
  • iddqd は、値からキーを借用する Rust のマップライブラリであり、Oxide の Omicron 制御プレーンでディスクや sled inventory のような大きなレコードのインメモリインデックスを維持するために使われており、正確性が重要
  • 標準の BTreeMap はキーと値を別々に保存するため受け渡しが煩雑になったり、重複したキーが食い違う可能性があるが、IdOrdMap はレコード内のフィールドからキーを抽出して参照する
  • unsafe Rust は、コンパイラが証明できない安全なプログラムを表現するための抜け道であり、ジェネリックコードがユーザー提供の trait 実装を呼び出すときは、病的な安全 Rust に対してさえ耐えなければならない
  • iddqd の mutable iteration は、インデックスがすべて異なるという不変条件に依存してライフタイムを延長しており、病的な Ord 実装によって B-tree と item set が食い違い、同じ項目に対する重複インデックスが作られる可能性があった
  • 修正ではキーとインデックスを一緒に比較し、失敗時にはユーザーコードを呼ばない線形スキャンにフォールバックするようにし、十分な信頼を得るには Miri・モデルベーステスト・panic fault injection・LLM による敵対的レビューを組み合わせる必要がある

iddqd が解決する問題

  • iddqd は、キーを値から借用するマップを提供する Rust ライブラリであり、Oxide はこれを Omicron 制御プレーンで広く利用している
  • Omicron は Oxide rack の中心で compute や storage のようなリソースをプロビジョニングし、rack が継続して動作するよう維持するソフトウェアであり、iddqd が誤動作すると制御プレーンが予測しづらく診断しにくい形で誤作動する可能性がある
  • Rust 標準ライブラリの BTreeMap<Email, User> のような構造は、キーである email を値とは別に保存する
  • キーと値を一緒に渡すには get_key_value で (email, user) を取得する必要があり、fetch 時点で UserRecord のような別構造体を作ると、レコード型が多い場合に扱いづらくなる
  • email を User の中にも複製すると、マップのキーと値の中の email が一致しなくなる危険が生じる
  • IdOrdMap は IdOrdItem trait によってレコードからキーを抽出させ、キー型は type Key<'a> = &'a Email のように値から借用できる
  • Oxide ではこのパターンがデータベース問い合わせ結果のような大きなレコードに適しており、制御プレーン全体での大きなレコードのインデックス作成に有用に使われている

追加機能

  • iddqd は、複数フィールドから借用する複雑なキーを第一級でサポートしており、dynamic dispatch のような回避策なしに扱える
  • BiHashMap と TriHashMap は、1つの項目をそれぞれ2つまたは3つのキーでインデックスし、複数のマップを手作業で維持する一般的なパターンを避けられる
  • Serde 実装はマップではなくシーケンスとしてシリアライズすることで、JSON で non-string key もシリアライズできるようにし、重複キーを拒否する
  • 下位互換性のため、マップ形式のシリアライズもサポートしている

unsafe Rust が難しくなるポイント

  • Rust における中核的な危険は未定義動作 (undefined behavior, UB) であり、安全なコードがどのような形でも UB を引き起こせないなら sound、そうでなければ unsound である
  • safe Rust ではコンパイラは UB を含むプログラムを拒否するが、静的解析の限界により、UB を含まない一部のプログラムもあわせて拒否してしまう
  • unsafe キーワードはそのようなプログラムを表現するための抜け道であり、書き手が soundness に対する責任を負ってコンパイラに信頼を求める仕組みである
  • unsafe コードが守るべき Rust の規則には、データ競合の禁止、未初期化または解放済みメモリの読み取り禁止、同じメモリ領域への複数の &mut alias の禁止、不変データの変更禁止がある
  • これらの規則は、特に mutable aliasing のために推論が難しく、通常は safe abstraction の背後に unsafe をカプセル化する必要がある

安全な抽象化を検証する段階

  • split_at_mut は mutable slice を2つに分割する安全なメソッドだが、safe Rust だけではこのような分割を表現できず unsafe が必要になる
  • split_at_mut の soundness は、&mut [T] を受け取っているか、mid <= slice.len() を確認しているか、残りの長さを正確に計算しているかといった周辺の安全コードの不変条件に依存する
  • MyVec<T> の get は、len が正確で範囲内であるという条件に依存しており、この条件は同じモジュール内の他のすべてのメソッドが維持しなければならない
  • ジェネリックな unsafe コードがユーザー提供コードを呼び出すと、難易度は最も高くなる
  • safe Rust のコードは、どれほど病的または敵対的に書かれていても、unsafe コードに UB を起こさせてはならない
  • ExactSizeIterator の len() を信じて capacity 分だけ書き込む collect_exact のようなコードは、iterator が偽の長さを返すと未割り当てメモリに書き込むことになり、一般には unsound である
  • std::io::Read も、バグのある実装が読み込んだバイト数を誤って返せる trait であり、Rust RFC 2930 がこの問題を扱っている
  • iddqd はユーザー提供の trait 実装を呼び出すジェネリックなデータ構造であるため、この最も難しいカテゴリに属する

iddqd の内部構造

  • iddqd は、項目の一覧である ItemSet と、slot インデックスを格納するテーブルを組み合わせている
  • IdHashMap<T> は、ItemSet<T> と ItemIndex を格納する hashbrown::HashTable を使う
  • ItemSet<T> は Vec<ItemSlot<T>> を持ち、ItemSlot<T> は Occupied(T) または Vacant { next: ItemIndex } である
  • free_head は、直近で解放された Vacant slot、または空き slot がないことを示す sentinel を指し、free_head と Vacant slot 群が free chain を形成する
  • 新しい項目を挿入するときは、free_head で再利用可能な slot を確認し、あれば Vacant slot を Occupied で上書きしたうえで free_head を次の値へ進める
  • 空き slot がなければ末尾に新しい Occupied slot を push し、その後キーを取得してハッシュを計算し、新しいインデックスを hash table に記録する
  • 削除は逆に、hash table でキーからインデックスを見つけて削除したあと、対応する ItemSlot を Vacant に変え、従来の free_head を next としてつなぐ
  • IdOrdMap は hash table の代わりに B-tree index を使い、BiHashMap と TriHashMap はそれぞれ2つと3つの hash table を保持する

mutable iteration とライフタイム延長

  • IdOrdMap::iter_mut は、キー順で項目を mutable iteration する
  • Rust の iterator は、返された値が iterator 自体を借用していてはならず、collect のようなコンビネータは iterator が消えたあとにも Vec<&mut T> のような値を残しうる
  • この動作が安全であるためには、iterator が同じ値を2回返してはならない
  • borrow checker は一覧への連続アクセスしか見ておらず、インデックスがすべて異なるという事実は認識できない
  • iddqd は std::mem::transmute::<&mut T, &'a mut T> によるライフタイム延長 (lifetime extension) を使っており、これは self.iter が返すインデックスがすべて異なるという不変条件に依存している

病的な Ord 実装が生んだバグ

  • キー 0 から 4 までの5項目が順番に入っている IdOrdMap で、Entry API によりキー 0 を検索すると、OccupiedEntry は内部に index 0 を保存する
  • その後 Key の Ord 実装が、値に関係なく常に Equal を返すように変わると、OccupiedEntry::remove が B-tree を再走査するとき、キー比較だけで誤った項目を削除してしまう可能性がある
  • たとえば B-tree でまず (key = 2, index = 2) を比較すると、Equal のためその entry が削除され、OccupiedEntry が保持していた index 0 により item set 側では item 0 が削除される
  • この状態では B-tree に index 0 が残っている一方で、item set の slot 0 は vacant となり、item 2 は item set に残っていても B-tree からのポインタを失う
  • その後 Ord が正常動作に戻り、キー 1000 の項目を挿入すると、free_head が指す slot 0 が再利用される
  • その結果、B-tree 内に同じ slot を指す重複インデックスが生じ、IterMut が同じメモリに対する複数の &mut 参照を作れてしまい unsound になる

修正方法と性能上のトレードオフ

  • B-tree をたどる際、キーだけでなく index の equality もあわせて確認するように変更した
  • 既知の index を持つキーを探す際、(key, index) の両方を比較するため、病的な Ord が Equal を返しても (key = 2, index = 2) と探している index 0 の比較は index の tie-breaker により Less になる
  • この検索が成功するには、保存されている index が探している index と実際に一致していなければならない
  • tie-breaker は誤った項目とのマッチを防ぐが、正しい項目を常に見つけられることまでは保証しない
  • B-tree はキーでソートされており、tie-breaker は index を比較するため、この2つの順序は一般に独立している
  • tree search が失敗した場合は、ユーザーコードを呼び出さない linear scan にフォールバックして B-tree から該当 index を削除する
  • この fallback により削除操作は logarithmic time ではなく linear time になるが、到達するのはバグのあるユーザーコードがある場合だけなので、受け入れ可能なトレードオフとされる

検証レイヤー

  • iddqd は基盤的なデータ構造として使われるため、分析的レビューと複数の実証的検証を組み合わせている
  • すべての unsafe block と unsafe pattern は、1人から3人の Rust 開発者およびレビュアーが分析する
  • 各 unsafe block の上には // SAFETY: コメントで reasoning を残し、Clippy の undocumented_unsafe_blocks lint でコメントの存在を確認している
  • example-based test はマップを作成して操作を行い、その結果を確認するもので、iddqd には内部 unit test と public API integration test がある
  • これらのテストは doctest としても存在し、ドキュメントの役割も兼ねている
  • 病的テストでは、バグのある Ord やその他の trait 実装を供給する
  • CI では regular test と pathological test を Miri の下で実行し、複数種類の UB を検出する
  • 重複インデックスが存在してはならないという不変条件のような条件は、Miri 外の通常のテスト環境でも確認できる

モデルベーステストと fault injection

  • iddqd は2層の randomized testing を使う。1つは正しい oracle と比較する model-based testing、もう1つはその上に重ねる fault injection である
  • model-based testing または stateful property-based testing は、ランダムな操作列を型インスタンスに適用し、結果を正しいと知られている oracle と比較する
  • iddqd は、非効率だが明らかに正しい NaiveMap oracle を基準に、広範な model-based test を実行している
  • fault injection は user code にバグをランダムに注入する方法であり、iddqd では panic safety が特に有効な軸だった
  • user code が操作の途中で panic しても、マップの不変条件が壊れてはならない
  • 各マップ操作にランダムな panic countdown を付け、user code の呼び出しごとに countdown を 1 ずつ減らし、0 になったら panic することでランダム panic safety testing を行う
  • この方法は iddqd における微妙なバグをいくつも発見し、その一部は unsoundness につながっていた
  • model-based test は各操作のあとに no-duplicate-index のような内部不変条件も検証する
  • model-based test は Miri の下で実行するには遅すぎるため、soundness と correctness が依存する不変条件を別途確認する

LLM による敵対的レビューと形式的手法

  • 現世代の frontier model は、マップを破壊する病的な user code 実装をいくつも見つけ出した
  • ある notable case では、custom allocator が panic して unwind するときにマップ不変条件が壊れる経路を LLM が構成した
  • これは、既存の panic safety test が扱っていた Ord 実装のような一般的な user code panic とは異なる panic safety 問題だった
  • LLM はもっともらしいが誤った soundness claim も作りうるため、防御策として Miri を oracle にした red-green TDD が有効になる
  • soundness bug では、まず Miri の下で UB を示すテストを作り、修正後に同じテストが通ることを再確認する
  • Kani のような model checker でマップ不変条件を証明するアプローチでは、iddqd は複雑すぎて必要な proof がツールの扱える範囲を超えるほど大きくなる
  • Creusot は Rust コードが panic やその他のエラーから自由であることを証明する助けにはなるが、user code が panic したり誤動作したりしても維持されるべき不変条件は現時点では証明できない
  • NaiveMap は iddqd に最も近い specification の役割を果たしており、CI で model-based test を数千回実行することで、実装が specification と一致していることへの高い信頼を得ている
  • cargo-anneal は、unsafe Rust の横にある doc comment に Lean の soundness proof を埋め込める開発中ツールとして関心を集めている
  • iddqd は、明確な不変条件と限定的ではあるが自明ではないスコープを持つため、形式検証ツールの benchmark 候補となる

結論

  • unsafe なジェネリック Rust は、任意の trait 実装や敵対的な実装まで考慮して各不変条件を維持しなければならないため、きわめて難しい
  • iddqd のバグは、病的な Ord 実装がマップを欺いて同じメモリに対する mutable alias を作れることを示している
  • 単一の手法だけですべてのバグを捕まえることはできないため、人間による line-by-line の unsafe reasoning、example-based test、pathological test、randomized test、frontier model による敵対的レビューを組み合わせている
  • Oxide は、基盤インフラにおいてはこのレベルの厳密さが正当化されると考えている

関連記事

1件のコメント

 
GN⁺ 9 시간 전
Lobste.rsのコメント

  • ちゃんと理解できていればですが（移動中にスマホで見ました）、これはラッパー型と、HashMap/BTreeMapの代わりにHashSet/BTreeSetを使えば解決できそうです
    ラッパー型は必須ではありませんが、安全性と今後の保守性を考えると良い選択です
    必要なのはオブジェクトを包むサイズ0のラッパーだけで、関心のあるフィールドだけを見るカスタムPartialEq/Hash実装を入れれば済みます。値全体を構築せずに検索したいなら、値型に対してAsRefを実装する2つ目の型を作ることもできます
    この方式は、unsafeなしで既存のHashSet/BTreeSetインターフェースをそのまま再利用する方法です。値/キー型を包む代わりに、こうした振る舞いを提供する新しいHashSet/BTreeSetラッパーを作ることもできます

    • ここでのキーはレコード型のフィールドやサブフィールドを任意に組み合わせたもので、GATで表現されています。また、整数インデックス/スラブのパターンは多重インデックスマップにも自然に一般化されます
      Entry APIや可変アクセス/反復なども含まれています。iddqdの中では可変性をかなり慎重に扱っていて、一部の箇所では内部可変性を許すためにアトミックな値を使っています。こうした処理はインデックス間接参照なしではかなり難しかったはずです

  • iddqdを形式検証するなら、最初はKaniのようなモデルチェッカーでマップが内部不変条件を壊さないことを証明しようとすると思います。ただ、iddqdはKaniが扱うには少し複雑で、必要な証明がツールの扱える規模を超えてしまうという話が気になりました
    この点をもう少し詳しく共有してもらえますか？ アルゴリズムの計算量が最悪ケースで劣化する、という意味なのか気になります
    全体として形式手法のケーススタディとして興味深く、この点に触れてくれたのはよかったです。既存の形式検証ツールについては、大規模プロジェクトでの成功例を見ると、少なくともデータ構造の正しさくらいは証明できるだろうと素朴に考えがちですが、実際にはデータ構造ごとに難しさが異なり、Rustのように無制限のエイリアシングを許す言語より証明に有利だと考えられがちな言語であっても、実務上は容易ではないことが分かります
    少し話はそれますが、図をどう作ったのかも気になります。一回限りのスクリプトを書いたのか、Oxideのブログ/ウェブサイトのデザインに合わせた特注品のように見えて、汎用ツールには見えません

    • 記事の末尾に「Diagrams courtesy Ben Leonard.」とあり、Ben LeonardはOxideのデザイナーです。なので、おそらく手作業の図だった可能性があります
    • 具体的でちゃんと動くtrait実装を対象に、ごく基本的なことを証明しようとしただけでも、CBMCはCPUを回し続けて10分以上経っても終わりませんでした
      範囲を狭めることも試しました。たとえばhashbrownが正しいと仮定しましたが、それでも大きな進展はありませんでした。そのあたりでほぼ諦めました。ちゃんと動くtrait実装についてはiddqdが正しいというかなりの確信がありますし、形式手法で本当に関心があったのは、任意の、壊れた実装に対しても成り立つ証明でした
      ただ、私はKaniの最も熟練した使い手というわけではありません。あなたや他の誰かがこれを試してくれたら本当にうれしいです
      図については、最初にMermaidでスケッチを作り、その後、素晴らしいデザイナーであるBen Leonardが、私たちのカラーパレットとテーマに合うよう手作業で仕上げてくれました

  • オブジェクトをそのフィールドの1つをキーにしてインデックスするフィールドベースのマップというパターンは、C#でもいつでも使えたらいいのにと常々思っています
    以前、自分で簡単なものを作ろうとしたことがありますが、インターフェースがあまりきれいにならず、結局やめました。この記事を見て、もう一度挑戦してみたくなりました

    • そうですね、本当に便利なパターンです。会社で必要な仕事のために作りましたが、その後はcargo-nextestのようなプロダクションコードベースから個人プロジェクトまで、あらゆるところで使うようになりました
      フィールド1つだけをキーにするのが最も一般的な使い方ではありますが、iddqdは十分に柔軟で、フィールド、サブフィールド、あるいはフィールドから純粋かつ低コストで計算できる関数の任意の組み合わせを使えます。たとえば https://docs.rs/iddqd/latest/iddqd/ でArtifactKeyを検索してみてください（Rustの構文に慣れていなければご容赦ください）
      iddqdを設計するとき、ユーザーがRust型システムの力をすべて使えるようにすべきだと強く考えました。ライブラリ作者である私がどれだけ多くの回り道をしなければならなくても構わない、という意味です。iddqdがユーザー、特に私の同僚たちにとって使っていて楽しいライブラリになってほしいと本当に思っていました