AI Agent向けSandbox実装

はじめましての方も、お久しぶりの方もこんにちは。スパイスコード代表の中河です。 スパイスコードは、「ロカルメ・オーダー」

order.localmet.com

というAI Agentを内包したERPサービスを開発・提供しているスタートアップです。 近年、AI Agentという言葉を耳にする機会が増えましたが、実は私たちはこのブームが来る前から、AI Agentの実用化に向けた開発に取り組んできました(例えば、本日ご紹介する機能の実装を行っていたのは2023年12月〜2024年1月頃です)。そして現在では、ERPの中核機能として、AI Agentを実際にお客様にご利用いただいています。 本ブログでは、私たちのAI Agentがどのような仕組みで動いているのか、そして他のAgentとは何が異なるのかについて、ご紹介していく予定です。 第1回目となる今回は、「AI Agentが生成したコードを安全に実行する独自のSandbox機構」について、その概要をご紹介します。

Agent Architecture

TL;DR: ロカルメ・オーダーでは ptrace を使った独自のSandbox機構を運用しています。本稿は短いサンプルコードでどのようにサンドボックスを構築しているのか?をご紹介いたします。

1. なぜいま “軽量サンドボックス” なのか

  • Agent 時代の要請: LLM エージェントは外部プラグインや自己生成コードを即座に実行する─安全な実行環境が必須。

  • 実行情報のフィードバック: ptrace で取得したシステムコールやリソース統計をリアルタイムに Agent へ返送し、自己改善するループを構築できる。

  • 実行環境の柔軟性: ptrace ベースはイメージ再ビルド不要でポリシーをコードから動的変更可能。オンプレや制約の厳しい閉域ネットワーク環境でも、Python 実行権さえあればそのまま導入できる。CI・サーバレス・ローカルデバッグまで同一コードで再利用可。

1.1 類似実装・活用例

ツール/プロジェクト 主な用途 ptrace の役割
gVisor (runsc, 旧 ptrace モード) コンテナ向けユーザ空間カーネル すべての syscall を ptrace でフックし、Go 実装の仮想カーネルで処理
LangChain Sandbox Agent向けsandbox Python codeをWebAssemblyにコンパイルして実行

2. ptraceとは?

ptrace() は トレーサ (親) がトレース対象プロセス (子) を観察・制御できる LinuxUNIX 系 OS のシステムコールです。レジスタやメモリの読み書き、シグナル挿入、システムコール前後での割り込み停止などを行えるため、デバッガ (gdb)、システムコールトレーサ (strace)、軽量サンドボックス、故障解析ツールなどの基盤として利用されています

2.1 主要リクエストと用途(代表例)

区分 リクエスト (req) 典型的な用途・ポイント
トレース開始 PTRACE_TRACEME 子→親へ「自分をトレースしてほしい」と通知。execve() 直後に gdb などが使う入口。
アタッチ/デタッチ PTRACE_ATTACH / PTRACE_DETACH 既存プロセスの動的アタッチと解除。解除時はシグナルや PTRACE_O_EXITKILL で安全に終了させることも可能。
停止制御 PTRACE_CONT / PTRACE_SYSCALL / PTRACE_SINGLESTEP 連続実行・システムコール毎停止・命令単位ステップ実行。監査・改竄ポイントとして SYSCALL 前後 で介入できる。
レジスタ I/O PTRACE_GETREGS / PTRACE_SETREGSx86 以外は PTRACE_GETREGSET 戻り値を -EPERM に書き換える、引数を編集するなど動的パッチを実施。
メモリ I/O PTRACE_PEEKDATA / PTRACE_POKEDATA, PTRACE_PEEKUSER ダンプ取得、コードインジェクション、データ改竄。
イベント通知 PTRACE_O_TRACECLONE TRACEFORK TRACEEXEC TRACESECCOMP など clone(2) / fork(2) / execve(2) / seccomp 発火時点で自動停止し、親がハンドリングできる。
安全終了 PTRACE_O_EXITKILL トレーサ異常終了時に tracee を自動 SIGKILLstrace --kill-on-exit が内部で利用 (man7.org)。

2.2 典型的なイベントループ (高水準フロー)

  1. waitpid(-1, &status, __WALL) で子の停止を待つ。

  2. WIFSTOPPED(status) のときかつ WSTOPSIG(status) が SIGTRAP+0x80 (PTRACE_O_TRACESYSGOOD) → システムコール入口 or 退出、同じくWIFSTOPPED(status) のときかつ status>>16 に PTRACE_EVENT_* が入っていれば PTRACE_GETEVENTMSG で追加情報を取得。

  3. システムコール入口なら PTRACE_GETREGS で orig_rax (x86-64)、第一引数(rdi など) を読み取る。

  4. ホワイトリスト外なら 戻り値を改竄 (PTRACE_SETREGS) したり PTRACE_KILL / SIGKILL を送る。

  5. 次の停止条件に合わせて PTRACE_SYSCALL または PTRACE_CONT で再開。

性能コスト: system call ごとに 2 回 ユーザ空間⇆カーネルを往復。I/O 多用処理は体感で~3× 遅れる。

なぜseccompを使用しないのか?: 単純なsystem callの許可/不許可だけではなく、Agent向けの他の処理も行っているからです(ここは別の機会にご紹介..)


3. Sandboxアーキテクチャ・実装

┌ sandbox.py (parent tracer)
│  ├─ fork → child (untrusted code)
│  ├─ waitSignals ⇆ syscall enter/exit
│  └─ timeout / rlimit / decision engine
└─→ 親が子の終了コードを返して終了

構成としては、コードを実行する worker プロセスと、その挙動を監視する親プロセス(Tracerプロセス)の二層に分かれています。Tracerプロセスでは、syscall(システムコール)の呼び出しを検知・制御するほか、必要に応じた各種セキュリティ処理を担っています。

3.1 実装(抜粋)

実装の主要な部分を抜粋してご紹介します。 本Sandboxはpython-ptrace https://pypi.org/project/python-ptrace/ を使用したPure Pythonで実装されており、対象のPythonコードを exec(3) を使わずに実行できる構造になっています。この仕組みにより、対象コードが利用するライブラリや機械学習モデルなどを、事前にロードしておくことが可能となり、より効率的かつ柔軟な実行環境を実現しています。

  def _spawn_sandbox_worker(self, server: socket.socket) -> None:
      pid, r, w = None, None, None
      try:
          r, w = os.pipe()
          pid = os.fork()   # MLモデル等を使用する場合はCoWを効かせる為、forkより前にロードする

          if pid == 0:
              try:
                  self._sync_child(r, w)
                  r, w = None, None

                  self._start_sandbox_worker(server)
              finally:
                  os._exit(0)
      finally:
          self._sync_parent(pid, r, w)

      self._logger.debug(f"Spawned sandbox worker {pid}")

  def _watch_sandbox_worker(self, pid: int, status: int) -> None:
      if os.WIFSTOPPED(status) and os.WSTOPSIG(status) == signal.SIGTRAP:
          reg = self._get_registers(pid)
          nr = get_syscall_number(reg)

          if not is_approved_syscall(pid, nr, reg):   # syscallが許可されたものか判断
              self._logger.warning(f"Unapproved syscall PID:{pid},SYSCALL_NR:{nr}")
              ptrace_kill(pid)
              return

      ptrace_syscall(pid)

  def _get_registers(self, pid: int) -> linux_reg | freebsd_reg | openbsd_reg:
      if HAS_PTRACE_GETREGS or HAS_PTRACE_GETREGSET:
          return ptrace_getregs(pid)

      words: List[Any] = []
      num_words = sizeof(ptrace_registers_t) // CPU_WORD_SIZE
      for offset in range(num_words):
          word: Any = ptrace_peekuser(pid, offset * CPU_WORD_SIZE)
          bytes = word2bytes(word)
          words.append(bytes)

      bytes = "".join(words)
      return bytes2type(bytes, ptrace_registers_t)

...

def get_syscall_number(
    reg: linux_reg | freebsd_reg | openbsd_reg,
) -> Any:
    if platform.machine() == "x86_64":
        return reg.orig_rax
    elif platform.machine() == "aarch64":   # docker on mac 環境用
        return reg.r8
    else:
        raise NotImplementedError


def is_approved_syscall(pid: int, nr: int, reg: linux_reg | freebsd_reg | openbsd_reg) -> bool:
    if platform.machine() == "x86_64":
        name = X86_64_SYSCALL_WHITE_LIST.get(nr)
    elif platform.machine() == "aarch64":
        name = AARCH64_SYSCALL_WHITE_LIST.get(nr)
    else:
        raise RuntimeError("Unsupported platform")

    if name is None:
        return False

    ...  # プロダクトではここでsystem callのgateway処理などを行っている

    return confirm_func(pid, reg)

実装上の工夫として、共通で使用され、かつメモリ消費の大きい機械学習モデルや大規模データは、fork前にロードする設計を採用しています。 これにより、Copy-on-Write(CoW)機構が有効に働き、子プロセス側で不要なメモリ確保が発生せず、全体として効率的なメモリ利用が可能になります。


4. プロダクトでの構成例

実際のプロダクトでは、Sandbox本体は「Server」と、Agentがツールとして利用する「Client」に分かれた構成になっています。 両者の通信は Unix Domain Socket を用いており、Sandbox自体はKubernetes環境上でSidecarとしてデプロイされています。


5. まとめと展望

  • 本手法のメリット: ユーザランドPython実装だけでAgent用の隔離環境が実現できる
  • 課題: ptrace オーバーヘッド、カーネル exploit への脆弱性、安全性を適切担保する為にはsystem call実装の理解が必要
  • 次の一手: 今回は、Sandbox機能にフォーカスしてご紹介しました。次回は、Agentの再現性を高めるために、処理系からどのような実行時情報をAgent側にフィードバックしているのかについて、ご説明できればと考えています。 ※他のコンテンツとの兼ね合いで、順番が前後する可能性があります;-)

参考文献