【ITニュース解説】A .nemo checkpoint is a tar file: two risks your model scanner may not open

NVIDIA NeMoは、最先端のAIモデルを開発・利用するための強力なフレームワークである。NeMoで訓練されたモデルは通常、.nemoという拡張子のファイルとして保存される。このファイルは一見すると、中身が不明な一つの大きなデータ塊のように思えるかもしれない。しかし、実際には.nemoファイルは、複数のファイルを一つにまとめたtarアーカイブ、つまり「圧縮されていない（あるいは古いバージョンではgzipなどで圧縮された）ファイルコンテナ」に過ぎない。このアーカイブの中には、モデルの設定情報が書かれたmodel_config.yamlファイル、モデルの学習済みデータであるmodel_weights.ckptファイル、そして時にはテキスト処理に必要なトークナイザーや語彙データなどがバンドルされている。NeMoがrestore_fromという機能を使ってモデルを復元する際、まずこの.nemoファイルを展開し、その中から必要なファイルをロードする仕組みになっている。

このような内部構造を持つ.nemoファイルだが、その性質を理解していないと、重大なセキュリティリスクを見過ごす可能性がある。多くの既存のモデルスキャナーは、.nemoファイルを単一の大きなファイルとして扱い、その中身であるtarアーカイブを展開してまで検査しないため、内部に隠された脅威を検出できないのである。これは、荷物の外側だけを見て中身の危険物を見逃してしまうような状況に似ている。具体的には、主に二つの大きなリスクが指摘されている。

一つ目のリスクは、「ピクル」と呼ばれるPythonのデータ形式に起因する脆弱性である。.nemoファイルの中に含まれるmodel_weights.ckptは、PyTorchのモデル保存形式であり、Pythonオブジェクトをバイト列に変換して保存する「ピクル」という仕組みを利用している。PyTorch 1.6以降では、zip形式の中にdata.pklというピクルファイルが含まれるが、それ以前の形式では直接ピクルデータが使われる。このピクル形式には、オブジェクトが復元される際に任意のPythonコードを実行できる「__reduce__ガジェット」という機能が隠されている場合がある。悪意のある.nemoファイルがこのガジェットを仕込んでいれば、モデルが復元される際に、スキャナーが予期しないコードが実行されてしまう恐れがあるのだ。一般的なスキャナーは.nemoファイルをtarアーカイブとして認識し、その中まで踏み込んでピクルの中身を詳細に分析する機能を持たないため、この種の危険なガジェットを検出できない。これは、近年研究で指摘されている、ピクルを悪用した攻撃経路の一つとして認識されている。

二つ目のリスクは、過去のNeMoバージョンに存在した「任意のファイル書き込み」の脆弱性（CVE-2024-0129）である。NeMoのバージョンr2.0.0rc0より前のリリースでは、_unpack_nemo_fileという関数がtar.extractall()メソッドを呼び出して.nemoファイルを展開していたが、この際にアーカイブ内に含まれるファイル名やパスが適切に検証されていなかった。例えば、アーカイブ内に../../somethingのような親ディレクトリを指定するパスや、絶対パスを持つファイルが含まれていた場合、モデルを復元する際に、本来展開されるべきディレクトリの外側にファイルが書き込まれてしまう可能性があった。これは、「CWE-22」というカテゴリに分類されるディレクトリトラバーサルの脆弱性の一種であり、攻撃者はこの仕組みを悪用して、システムの重要なファイルを上書きしたり、新たなファイルを書き込んだりすることができる。これにより、最悪の場合、システムのコード実行権限を奪われる危険性も指摘されており、セキュリティ評価機関では高い危険度が設定されていた。現在も古いNeMoバージョンを利用している環境では、信頼できない.nemoファイルをロードするだけで、意図しないファイルがシステムに書き込まれる可能性がある。

これらのリスクに対処するため、keras-model-scanというオープンソースのツールキットにnemo_container_scan.pyという新しいスキャンツールが追加された。このツールは、restore_fromが.nemoファイルを処理するのと同様の方法で、内部構造を安全に検査することを目指している。具体的には、.nemoファイルをディスクに展開することなく、メモリ上で標準ライブラリの機能を使ってtarアーカイブを開き、その中にあるメンバー（ファイル）を一つずつ順番に走査する。もしネストされたtorch.saveのzipファイルが存在すれば、それも剥がして内部のピクルファイルをpickletoolsを使って逆アセンブルし、その中身を詳細に分析する。この際、PyTorchやNeMoのライブラリを実際にインポートしたり、extractallのような展開関数を呼び出したり、ピクルを実際に復元・実行したりすることは一切ない。これにより、悪意のあるコードがスキャン中に実行されるリスクを排除しつつ、内部に潜む危険な要素を検出できるのである。ツールの設計では、ファイルサイズやメンバー数にも上限を設けることで、過剰なリソース消費や悪意のあるアーカイブによる攻撃にも耐えうるようになっている。

このツールは、以下のような特定のセキュリティルールに基づいて潜在的な問題を検出する。例えば、「NEMO001」は、コンテナ内のピクルがコード実行を引き起こす可能性のある呼び出し可能オブジェクトをインポートしている場合に「高」の重大度で警告を発する。また、「NEMO010」は、危険なインポートはないものの、モデル復元時にアンピクルされるピクルチェックポイントが含まれていることを「中」の重大度で示す。これは通常のNeMoチェックポイントでも発生する可能性があり、その場合は注意喚起として捉えるべきである。「NEMO020」は、tarメンバー名が抽出ルートから逸脱する可能性がある場合に「高」の重大度で、以前のCVE-2024-0129と同様の脆弱性を示す。さらに、「NEMO021」は、tarアーカイブ内のシンボリックリンクやハードリンクがアーカイブ外を指している場合に「中」の重大度で警告する。このツールは、何か「高」の重大度で問題が検出された場合に終了コード1を返すため、継続的インテグレーション（CI）の自動テストプロセスに組み込むことで、問題のある.nemoファイルがシステムに組み込まれることを防ぐためのチェックポイントとして利用できる。また、一部の.nemoファイルに見られるような、tarアーカイブの末尾に不適切なデータブロックがあっても、tar.next()メソッドを使用してメンバーを読み込むことで、エラーなく処理できる工夫も凝らされている。

ただし、このnemo_container_scan.pyツールは静的解析ツールであり、実際に脆弱性が悪用可能であることや、特定のファイルが必ずしも攻撃を引き起こすわけではないことを証明するものではない。あくまでもコードやデータの構造を分析し、潜在的なリスクの兆候をフラグ立てするものである。ツールの限界として、一度に解析できるネストされたzip層は一つまでであり、標準ライブラリが読み取れない特殊な形式のファイルは開けない。また、前述の「NEMO010」のように、通常の安全なNeMoチェックポイントに対しても警告を発することがあるが、これはピクルが持つ潜在的なリスクを意識するための注意喚起として受け止めるべきである。このツールの目的は、nemoファイルが単なる不透明なデータ塊ではなく、tarアーカイブであり、モデルのrestore_from機能がその中身を展開してロードする仕組みであるという本質的な理解を促すことにある。そして、その理解に基づいて、セキュリティスキャナーも同様に.nemoファイルの内部まで検査すべきだという重要なメッセージを伝えることにある。AIモデルの利用が拡大する中で、ファイル形式の深い理解とそれに基づいたセキュリティ対策がますます重要になっている。