Neural Network Inference with Intel SGX and Rust

第三者の計算機でニューラルネットワークを動かしたいが、第三者にニューラルネットワークを知られたくは無いという問題を想定して、その1つの解答である機密コンピューティング（Intel SGX）について、仕組みと簡単なPoCをご紹介します。

code: https://github.com/kokeshing/dnn-with-rust-sgx

Transcript

1. AI 2024.12.05 大西 弘太郎 株式会社ディー・エヌ・エー ＋ GO株式会社 持ち主に秘密で計算しよう！ ~AI in Intel SGX~
2. AI • 本スライドの内容は、細心の注意を払って作成しておりますが、 その正確性について保証するものではありません • 本スライドの内容は、2024年12月当時の情報を元に作成しています • 本スライドの内容には詳細・背景に言及できていない箇所もあるの で、末尾にある参考文献も合わせてご参照ください • 実用性はありません ◦ 実用性はありません ▪ 実用性はありません 2 おことわり
3. AI 3 • 大西 弘太郎 • 株式会社ディー・エヌ・エー • ソリューション本部 • データ統括部 • AI技術開発部 • 主なお仕事 自己紹介
4. AI 4 お品書き 01｜概要 02｜Trusted Execution Environment (TEE) とは 04｜Intel SGXでAIモデルを動かす 03｜Intel Software Guard Extensions (SGX) とは
5. AI 5 01 概要
6. AI 超高性能 or 独自技術 or 独自IPなAIモデルができた！ ↓ 通信量・レイテンシーを抑えたいからユーザーの手元で動かしたい！ ↓ ユーザーのデバイスだとメモリやストレージのAIモデルを見られちゃうぞ🤔 ↓ ユーザーにはAIモデルを秘密にしたまま計算したい😭 ということありませんか？（諦めてサーバーで処理しよう） ⇒ Confidential computing (機密コンピューティング) をしてみよう！ 　　　　　　 6 概要
7. AI 7 今回のゴール PC内のAIモデルを使える PC内のAIモデルがバレない
8. AI • 今回行おうとしていること（の関連） ◦ アンチリバースエンジニアリング ◦ Digital Rights Management (DRM) • リバースエンジニアは常にデバイスにフルアクセスすることが できるので、（リソースがあれば）必ず勝利することができる 👉 厳しい脅威モデルを想定する必要あり ハードウェアベースでデータを保護する技術である TEEの1つ、Intel SGXを 使ってセキュアにAIモデルを推論しよう！ 8 なぜConfidential computing？
9. AI 9 02 Trusted Execution Environment (TEE) とは
10. AI 10 Trusted Execution Environment: 信頼可能な実行環境 • ハードウェアを用いて実現されたセキュアな実行環境 • 信頼可能な領域 ・信頼しない領域 を定め、信頼可能な領域で データ・プログラムを扱うことで、セキュアな実行環境を実現 ◦ 信頼可能な領域 : 信頼可能なハードウェア・信頼可能な ハードウェアにより生成された保護領域 ◦ 信頼しない領域 : 上記以外全て • 他のハードウェア隔離実行技術（TPMなど）と異なり、一般の開発者 がTEEで動くプログラムを書くことができる
11. AI 11 TEEの実例 😄 様々なTEE実装が複数のベンダーから提供されている • Intel Software Guard Extensions (SGX) • Intel Trust Domain Extensions (TDM) • ARM TrustZone • RISC-V KeyStone • AMD Secure Encrypted Virtualization (SEV) 👆 メモリ暗号化する/しない、VM・メモリごと隔離する/しない など、隔離実行環境をどう提供するかも様々
12. AI 12 TEEの応用例 🤗 TEEは身近な場面でたくさん活用されている • PIN・指紋認証・顔認証などの鍵管理 ◦ e.g., Android Keystore・Secure Enclave (iOS) • デジタルコンテンツ管理（DRM） • 金融・決済 • クラウド上での機密コンピューティング
13. AI 13 TEEの応用例（機械学習） 🤖 最近は機械学習関連の話題でも • Apple: Apple Intelligenceで利用されるクラウド計算環境はTEEを利 用 • Google: モバイルデバイスとクラウド計算環境でLLMを動かす際に TEEを利用する実装を提示 👆 どちらも個人情報などをデバイスの外のクラウド上で推論する ための秘密計算用途 👿 秘密計算用途ではRA（後述）の致命的な問題が...
14. AI 14 TEEを使える環境（クラウド） 🤔 TEEを一般開発者が使う際の選択肢は？ • クラウドベンダー各社はTEEを用いた様々な 機密コンピューティングソリューションを提供 ◦ AWS: Nitro Enclaves ◦ Google Cloud : Confidential VMs, Confidential GKE Node ◦ Azure: Confidential VM 👉 コンシューマ向けの PCやスマートフォンでは？
15. AI 15 TEEを使える環境（ローカル） 😢 一般開発者が手元で気軽に使える TEEはまだ限られている • Android: Android用にTEEを動かすためのOSであるTrustyは、サー ドパーティーアプリから使うことはできない • iOS: 一般開発者はSecure Enclaveを利用した暗号化APIなどのみ 使うことができる • PC (ほぼLinux): プロセッサの搭載TEEが概ね利用可能 ◦ Intel SGX・TDM ◦ AMD SEV ◦ OP-TEE (Arm TrustZoneの制御実装の1つ）
16. AI 16 03 Intel Software Guard Extensions (SGX) とは
17. AI 17 Intel Software Guard Extensions (SGX) の概要 • Intel CPUに搭載されている専用ハードウェアを用いて、 メモリ上にEnclave (飛び地) と呼ばれる保護領域を生成し、 プログラムを隔離実行する • 搭載されているCPU ◦ Core 第6世代以降〜第10世代以前 ◦ Xeon SP 第2世代の一部 ・第3世代以降〜 👆 デスクトップ向けの CPUには搭載されなくなってしまった
18. AI 18 Intel SGXが想定する脅威モデル 出典: Chia-che Tsai, Donald E. Porter, Mona Vij. "Graphene-SGX: A Practical Library OS for Unmodified Applications on SGX." 2017 USENIX
19. AI 19 Intel SGXが想定する脅威モデル 🤯 BIOS・Host OS・Hypervisorすらも信頼していない • Intel SGXの信頼する領域は3つだけ ◦ CPU ◦ CPUにより生成されるメモリ上の保護領域 ◦ CPUとメモリ上の保護経路 • 同一マシンの他プロセスはもちろん、OSレベルでは特権を持つ 攻撃者、CPU以外のハードウェアへの攻撃も脅威として想定
20. AI 20 Intel SGXのメリット 👺 非常に厳しい脅威モデルによる、堅牢な隔離実行環境 ◦ たとえ特権があっても、Intel SGXには勝てない ◦ 他のTEEには、OS・Hypervisorについては信頼可能な領域 としているものもある 🚄 効率的にプログラムを動かすことが可能 ◦ TEEに共通するメリットで、ハードウェア実装されていること・計算 コストの小さい暗号化によって、速度・メモリ使用量の両方で高パ フォーマンスを期待できる ◦ CPUの機能をフルに使うことができるので、SIMDも使える
21. AI 21 Intel SGXのデメリット 👾 脆弱性はあるさ /発生するさ ◦ 日夜攻撃対象にされて脆弱性を見つけられている ◦ SGX自体が堅牢でも、SGXで動かすソフトウェアに脆弱性がある とダメ🙅 ▪ e.g., サイドチャネル攻撃・仕組みの不理解 👹 凶悪な実装難易度 ◦ 仕組みが難解 ◦ 公式SDK・ドキュメントが難解
22. AI 22 Intel SGXのデメリット 👹 凶悪な実装難易度 ◦ 公式のドキュメントがメンテされてない ◦ 周辺のライブラリ・ドキュメントもメンテされてない ◦ リソースや機能の制限がある ▪ OSを信頼しない ⇒ システムコール使えない ▪ Enclaveサイズに制限あり（96 MB or 512 GB） ◦ 環境を用意するのが大変（オンプレ・ローカル） ◦ ニッチすぎて情報がない
23. AI 23 04 Intel SGXでAIモデルを動かす
24. AI 24 今回のゴール（再掲） PC内のAIモデルを使える PC内のAIモデルがバレない
25. AI 25 はじめに ここからは実際に Intel SGXでAIモデルを動かす例をご紹介 • 以下の点も見ていきます ◦ Intel SGXの信頼性を支える仕組み ▪ TEEはただの隔離実行環境なので、信頼性を担保する 仕組みが必要 ◦ Intel SGXを使いやすくする営み ▪ 巨人の肩に乗ろう • 今回の例を実装したレポジトリは こちら
26. AI 26 Intel SGX環境の用意 • Intel SGXが動くマシン探しの旅 ◦ デスクトップ機（第7世代）⇒ ライブラリが動かず断念 ◦ サーバー（第10世代）⇒ チップセット側が対応せず ◦ サーバー（第9世代）⇒ 有効にするとなぜかGRUBで落ちる 😭 ローカル環境を用意するのは断念 👉 Azureが提供する Intel SGXが使える仮想マシンを使用 ◦ レポジトリに環境を用意するterraformを用意しています
27. AI 27 Enclaveプログラムの実装 • SDKを使ってC++を書いていけば実装できる...が ◦ 使用できる関数に制約がある ◦ Enclave内ではシステムコールが使えないので、Enclave外で動く プログラムも一緒に用意して解決させる ▪ Enclave境界を何度も反復横跳びする開発は非常に大変 ◦ 共有ライブラリはもちろんEnclaveプログラムには使えない ◦ SGXが堅牢なので、安全な言語で書きたい... 👉 Fortanix Rust Enclave Development Platform
28. AI 28 Fortanix Rust Enclave Development Platform (EDP) • Fortanix社が提供するIntel SGX向けのRustプラットフォーム • 専用ターゲットを指定 + 専用ランナーを使うことで、通常のRustコード をIntel SGXで動かせる ◦ cargo +nightly run --target x86_64-fortanix-unknown-sgx • 制約はあるものの、通常のRust開発とほぼ変わらない ◦ 一部の標準ライブラリが使えない ◦ Pure Rustでないライブラリが使えない 🤔 システムコールが使えないなどの制約をどう解決？
29. AI 29 Fortanix Rust EDPのアーキテクチャ 🙏 独自インターフェースを用意することで環境差異を吸収 • すべてのシステムコールが使えるわけではない（20未満） • 専用ターゲットによって標準ライブラリを対応
30. AI 30 RustとIntel SGX Fortanix Rust EDP以外にもRustでSGX開発が行えるツールは存在 • rust-sgx-sdk: SGX SDKのRustラッパー • SCONE: muslベースのC標準ライブラリインターフェースを 用意・それを用いる独自ターゲットも用意 ◦ TensorSCONE: SCONEを用いてTensorFlowでの学習・推論を Intel SGX上で行う論文
31. AI 31 AIモデルを動かす流れ 〜はじめに〜 ここからは大まかな流れをご紹介 • Intel SGXでAIモデルを動かすために3つのプログラムを用意 ◦ ISV-Client: リクエストを投げるクライアント（非Enclave） ◦ ISV-Enclave: リクエストを受け取ってAIモデルで推論するサー バー（Enclave） ◦ SP: AIモデルをEnclave内のプログラムに配布する外部 サーバー（非Enclave・別マシン） ▪ 〆切の都合上、 実際のプログラム では省略
32. AI 32 AIモデルを動かす流れ 〜概要〜 Intel SGXでAIモデルを動かす流れはシンプル 1. AIモデルを外部サーバからEnclaveで動作するサーバーに送信 ⚠ セキュリティ・ SGXの仕組み上、絶対必須 🙇 〆切の都合上、 実際のプログラム では省略 2. 推論リクエストをクライアントからEnclave動作サーバーに送る 3. 受け取ったリクエストを元にAIモデルを推論 4. 推論結果をEnclave動作サーバーからクライアントへ返却
33. AI 33 AIモデルを動かす流れ 〜AIモデルの配布〜 1. AIモデルを外部サーバーから Enclaveで動作するサーバーに送信 • 今回の要件: PC内のAIモデルがバレない 🤔 Enclaveを動かしているTEEは想定通りの隔離実行環境か？ 🤔 送信先のEnclaveは想定通りのプログラムか？ 🤔 AIモデルを別プロセスに傍受されないようにEnclaveへ送信でき るのか？ 👉 Remote Attestation（リモート構成証明）
34. AI 34 AIモデルを動かす流れ 〜Remote Attestation〜 TEEの信頼性を支える Remote Attestation • Intel SGXに限らず、TEE全般でTEE自体が正常なことを証明するこ とは必須 ◦ クラウド上で動くTEEがクラウドベンダから隔離されているのか？ ◦ ユーザーの端末のTEEがユーザーから隔離されているのか？ • 第三者機関（Intel）にマシンとEnclaveに対して唯一 となる レポートキー を検証を依頼することで、マシンとEnclaveの 真正性を証明
35. AI 35 AIモデルを動かす流れ 〜Remote Attestationの流れ〜 出典: Sakurai, A. (2023).【技術】Intel SGX Attestation詳説 - EPID Remote Attestation編. 株式会社Acompany. Retrieved December 25, 2024, from https://acompany.tech/privacytechlab/sgx-remote-attestation 詳細は参考文献を参照
36. AI 36 AIモデルを動かす流れ 〜信頼の基点 (Root of Trust)〜 🤔 Remote Attestationでどうやってマシンが改ざんされていない ことを確かめるのか 👉 信頼の基点 (Root of Trust) ◦ デバイスの信頼性を保証するためのコンポーネント ◦ Intel SGXの場合、Intel Management Engineとその中にある CPU固有の鍵 ▪ 鍵が漏洩しない限り、隔離された実行環境であるように設計 されている（= 信頼できる根拠） ◦ Remote Attestationでも対象の鍵が漏洩している場合は失敗
37. AI 37 AIモデルを動かす流れ 〜AIモデルの配布〜 1. AIモデルを外部サーバーから Enclaveで動作するサーバーに送信 1. 外部サーバー（SP）からEnclaveで動作する サーバー（ISV-Enclave）にRemote Attestationを行い、受理でき ることを確認 2. Remote Attestationに並行する形で、セッション共通鍵も 共有される 3. 共有鍵を用いて、外部サーバーとEnclaveで動作する サーバーの間で暗号通信路を確立 4. 暗号通信路を用いてセキュアにAIモデルを送信
38. AI 38 AIモデルを動かす流れ 〜Remote Attestationの恐怖〜 👿 Remote Attestationの流れは非常に難解・複雑（ 参考文献） • 仕組みを理解した上で受理判定などを実装しなければ致命的な セキュリティホールに • Fortanix Rust EDPでもRAのサンプルは最新のexamplesからは 消えている • Fortanix Rust EDPに頼らないRustでのRAのサンプルも、コード ベースが古くまず動かすところから難産 🙇 実際のプログラム に実装するのは間に合いませんでした
39. AI 39 AIモデルを動かす流れ 〜推論のリクエスト〜 2. 推論リクエストをクライアントから Enclave動作サーバーに送る 1. Enclave動作サーバーはTcpListenerでリクエストを待つ 2. クライアントからTcpStream::connectでコネクションを 確立 3. クライアントから推論してほしいデータをリクエストとして送付 👆 Fortanix Rust EDPではTCP通信ができるようになっている 👆 cargo workspaceを使ってクライアントとサーバー共通で使う リクエスト・レスポンス構造体などを共有すると楽
40. AI 40 AIモデルを動かす流れ 〜AIモデルの推論〜 3. 受け取ったリクエストを元に AIモデルを推論 1. 受け取ったデータを前処理 2. データをAIモデルへの入力として計算 3. AIモデルの推論結果からレスポンスを生成して返却 • 共有ライブラリや外部ライブラリが使えないので、ネットワークの計算 部分はRustで完結するように（今回はndarrayを使用） ◦ 全結合層や活性化関数の計算は手動で定義しておく ◦ モデルの計算も上記を使って手動で定義
41. AI 41 AIモデルを動かす流れ 〜AIモデルの推論〜 4. 推論結果を Enclave動作サーバーからクライアントへ返却 • TCP通信のレスポンスとして推論結果を返すだけ 以上で一連の流れは終了！
42. AI 42 まとめ
43. AI 43 まとめ • TEEについてご紹介 • TEEの１つであるIntel SGXについてご紹介 • Intel SGXを使ってAIモデルを推論させてみる例をご紹介 ◦ Intel SGXの信頼性を支える仕組みをご紹介 ▪ Remote Attestation ▪ 信頼の基点 ◦ Intel SGXを使いやすくする試みについてご紹介 ▪ Fortanix Rust Enclave Development Platform
44. AI 44 おまけ 〜実用性について〜 💁 A. 無し • 特定のCPUに依存 ◦ 4年前に出たCPUが家庭用としては最後... • Intel SGX自体が数々の問題を抱えている ◦ SGX.Fail（日本語解説）: SGXの実ケースにおける運用の 実態を調査、SGXが理想とする運用と現実は乖離しており、セ キュリティと使いやすさのトレードオフから結果的に セキュリティが毀損されていることを示す研究
45. AI 45 参考文献 日本語でIntel SGXの情報に容易にアクセスできることにこの場を借りて感謝申し上げます󰢛 1. aos. (n.d.).セキュリティキャンプ - 1 TEEとは何か？. Cliffford Terminal. Retrieved December 25, 2024, from https://qliphoth.io/media/2024-1-what-is-tee.pdf 2. aos. (n.d.).セキュリティキャンプ - 2 Intel SGXの基礎. Cliffford Terminal. Retrieved December 25, 2024, from https://qliphoth.io/media/2024-2-intel-sgx-basics.pdf 3. aos. (n.d.).セキュリティキャンプ - 6 EPID Remote Attestation. Cliffford Terminal. Retrieved December 25, 2024, from https://qliphoth.io/media/2024-6-epid-ra.pdf 4. aos. (n.d.).セキュリティキャンプ - 8 SGX Fail. Cliffford Terminal. Retrieved December 25, 2024, from https://qliphoth.io/media/2024-8-sgx-fail.pdf 5. Sakurai, A. (2022).【技術】TEE（Trusted Execution Environment）とは？. 株式会社Acompany. Retrieved December 25, 2024, from https://acompany.tech/privacytechlab/trusted-execution-environment 6. Sakurai, A. (2022).【技術】Intel SGXの秘密計算利用の難しさとは？. 株式会社Acompany. Retrieved December 25, 2024, from https://acompany.tech/privacytechlab/tee-intel-sgx 7. Sakurai, A. (2023).【技術】Intel SGX Attestation詳説 - Local Attestation編. 株式会社Acompany. Retrieved December 25, 2024, from https://acompany.tech/privacytechlab/sgx-local-attestation 8. Sakurai, A. (2023).【技術】Intel SGX Attestation詳説 - EPID Remote Attestation編. 株式会社Acompany. Retrieved December 25, 2024, from https://acompany.tech/privacytechlab/sgx-remote-attestation
46. AI 46 参考文献 9. Sakurai, A. (2023).【技術】Intel SGXの難しさとジレンマ. 株式会社Acompany. Retrieved December 25, 2024, from https://acompany.tech/privacytechlab/intelSGX-dilemma 10. 須崎 有康 (2020). Trusted Execution Environmentの実装とそれを支える技術. 電子情報通信学会 基礎・境界ソサ イエティ Fundamentals Review, 14(2), 107-117. Retrieved December 25, 2024, from https://typeset.io/pdf/implementation-of-trusted-execution-environment-and-its-4pyvkqgl67.pdf 11. 須崎 有康 (2021).TEE(Trusted Execution Environment)とそれに関する研究開発動向.デジタルサービス・プラット フォーム技術 特別研究専門委員会. Retrieved December 25, 2024, from https://www.ieice.org/~dpf/wp-content/uploads/2021/08/DFS%E7%A0%94%E7%A9%B6%E4%BC%9A%E7 %94%A3%E7%B7%8F%E7%A0%94%E9%A0%88%E5%B4%8E.pdf 12. Schaik, S., Seto, A., Yurek, T., Batori, A., AlBassam, B., Garman, C., Genkin, D., Miller, A., Ronen, E., & Yarom, Y. (2022). SoK: SGX.Fail: How Stuff Get eXposed. Retrieved December 25, 2024, from https://sgx.fail/ 13. Intel, Corp. (n.d.). インテル® ソフトウェア・ガード・エクステンションズ (インテル® SGX). Retrieved December 25, 2024, from https://www.intel.co.jp/content/www/jp/ja/products/docs/accelerator-engines/software-guard-extensions.html 14. Fortanix, Inc. (n.d.). Fortanix Enclave Development Platform . Retrieved December 25, 2024, from https://edp.fortanix.com/