RustでWebAssemblyインタプリタを作った話

RustでWebAssemblyインタプリタを作った話

本記事は「Wantedly 新卒 Advent Calendar 2021」の14日目の記事です。

5,6月にWebAssemblyインタプリタを作ったのでそれに関して話していきます。
主にWebAssemblyのフォーマットについての説明やどのような流れで開発を行ってきたかを話していきます。

コードはGitHubにて公開しているので興味のある人は見てみると面白いかもしれないです >_< (50 starちょうどで切りが良いのですがまだまだスターが欲しいですください)

https://github.com/k-nasa/wai

話さないこと

この記事ではWebAssemblyの概要に関する説明は行いません。WebAssemblyの目的や今後どのようなところで使われていくかなどは次のスライドにて丁寧に説明されているので興味のある人は見てみると良いと思います

https://www.slideshare.net/TakayaSaeki/webassemblyweb-244794176

Introduction

今回作ったインタプリタはwaiと名付けているのでここからはWebAssemblyインタプリタではなくwaiと呼ぶことにしましょう
では今回作ったwaiについて話していきます。

waiはRustで書かれたWebAssemblyの実行環境です。wasm形式のバイナリを入力として与えることでそれを実行し結果を返してくれます。
現段階で170命令の内100命令ほど実装している状態で割と動く形にはなってきています。（wasmの命令は続々と増えているのですでに170では無いかもしれませんが）

waiのデモとしてフィボナッチ数列の計算するwasmバイナリを与えた時の様子を貼っておきます。いい感じに動いてくれていますね。

WebAssembly Overview

まず最初はwasmのバイナリフォーマットを学ぶところからはじめました。
主に2つのサイトを参考にしながらフォーマットを完全に理解しました。1つはWebAssembly Specificationでもう一つは WebAssembly/designというリポジトリです。

WebAssembly Specificationはwasmの使用が全てまとまっているのでこれを読んでおけば詳細な情報を得ることが出来ます。動きがよく分からない命令があるときはspecを読んで試行錯誤しながら実装していました。

ただspecの文法の表記方法が見慣れないものだったのでバイナリフォーマットや命令の入出力を把握するのは難しく感じました。その時は2つめの WebAssembly/designを読みながら実装していました。
特にバイナリフォーマットがいい感じにまとまっていて、WebAssembly/designのBinaryEncoding.mdを見るとwasmのdecodeには困らなかったです。デコーダーを書くときはWebAssembly/designリポジトリを参照して、命令を解釈して実行していくときはWebAssembly Specificationを見るのが楽かもしれません。

waiでは次の2工程に分けてwasmbバイナリを実行しています。

1. バイナリのdecode
2. decodeした命令列を実行していく

wasmは次のようなバイト列になっています。これをプログラムで解釈しながら実行するのは困難なので、一度人間に優しい形式にデコードして実行しています。

:) % hexyl examples/add.wasm
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 00 61 73 6d 01 00 00 00 ┊ 01 07 01 60 02 7f 7f 01 │0asm•000┊•••`••••│
│00000010│ 7f 03 02 01 00 07 07 01 ┊ 03 61 64 64 00 00 0a 09 │••••0•••┊•add00__│
│00000020│ 01 07 00 20 00 20 01 6a ┊ 0b                      │••0 0 •j┊•       │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

↑のバイナリをWATというWebAssemblyのテキストフォーマットで書くと次のようになります。
2つの数値を足し合わせる関数を定義しています。

(module
  (type (;0;) (func (param i32 i32) (result i32)))
  (func $add (type 0) (param $lhs i32) (param $rhs i32) (result i32)
    local.get $lhs
    local.get $rhs
    i32.add)
  (export "add" (func $add))
)

decode

wasmは次のようなセクションが順番に並んでいる構造を取ります。

section number	section
0x00	custom section
0x01	type section
0x02	import section
0x03	function section
0x04	table section
0x05	memory section
0x06	global section
0x07	export section
0x08	element section
0x09	start section
0x0A	code section
0x0B	data section

このセクションはそれぞれ必須というわけではないので順番通りにデコードしていくことはできません。
各セクションの先頭にはsection typeを表すセクション番号とセクションの長さが入っているのでそれを見ながらsection typeとセクションに対応するバイトコードの解決を行っていきます。

セクションの長さとtypeの解決は次のようなコードで行っています。

pub(crate) fn decode_section_type(&mut self) -> Result<(SectionType, u32), DecodeError> {
    let mut section_number = [0; 1];

    // section numberは8byte固定なので8byte読む
    self.reader.read_exact(&mut section_number)?;
    let section_type = SectionType::from(section_number[0]);

    // sectionのsizeは可変長整数で指定されるのでいい感じに読む
    let section_size = self.decode_ver_uint_n()?;

    Ok((section_type, section_size.into()))
}

ここまでで得られたsection typeとsection sizeを使ってそれぞれのセクションをデコードしていきます。あとは愚直にデコードしていくだけですね。

pub(crate) fn decode_section(
    &mut self,
    section_type: SectionType,
    section_size: u32,
) -> Result<Section, DecodeError> {
    let section = match section_type {
        SectionType::Custom => self.decode_custom_section(section_size)?,
        SectionType::Type => self.decode_type_section(section_size)?,
        SectionType::Import => self.decode_import_section(section_size)?,
        SectionType::Function => self.decode_function_section(section_size)?,
        SectionType::Table => self.decode_table_section(section_size)?,
        SectionType::Memory => self.decode_memory_section(section_size)?,
        SectionType::Global => self.decode_global_section(section_size)?,
        SectionType::Export => self.decode_export_section(section_size)?,
        SectionType::Start => self.decode_start_section(section_size)?,
        SectionType::Element => self.decode_element_section(section_size)?,
        SectionType::Code => self.decode_code_section(section_size)?,
        SectionType::Data => self.decode_data_section(section_size)?,
        SectionType::Unsuport => todo!(),
    };

    Ok(section)
}

wasmはかなりシンプルな構造をしているためデコードは容易だと思います。
前述のWebAssembly/designにあるBinaryEncodingのドキュメントがかなり役立ちました。 https://github.com/WebAssembly/design/blob/main/BinaryEncoding.md

decoderのテスト

decoderをガリガリ実装した後はファジングテストにていい感じにデコードできるのか、クラッシュするwasm binaryはないかをテストしました。
このときのテストケースはwasm-smith/wasm_smithというcrateを使って生成しました。このcrateのおかげてたくさんのテストケースを書く必要がなかったので非常に感謝しています 🙏

cargo-fuzzとwasm-smith/wasm_smithを使うことでほとんど手を動かさずにテストを行うことができました。(テストが充実していると開発体験が非常に良いですね)

use wai::*;
use libfuzzer_sys::fuzz_target;
use wasm_smith::Module as M;

fuzz_target!(|module: M| {
    let bytes = module.to_bytes();

    Module::from_byte(bytes).unwrap();
});

fetch function

ここまででwasmモジュールのデコードが完了したのであとは一つずつ命令を実行していくだけになります。

waiのインターフェースは次のように関数名を受け取りそれを実行する用になっています。

wai examples/add.wasm --invoke add -a 1 2

デコードの段階で関数名とその実装をfunctio_tableという形で持っているのでそれを見て、引数で与えられた関数名から実行すべきコードの解決を行っています。
また関数の型も分かっているので引数の値や個数が実行すべき関数と合うかvalidationを行い、その後実行しています。

関数の返り値はstackに積まれているのでそこから必要な個数取り出して最終結果としています。(wasmの関数は複数の値を返す)

pub fn invoke(
    &self,
    name: impl AsRef<str>,
    args: Vec<RuntimeValue>,
) -> Result<ValueStack, RuntimeError> {
    let index = self.resolve_function_name(name.as_ref());

    let function_table = self.function_table();

    let func = function_table.get(index).unwrap();
    let return_length = func.returns.len();

    let memory = Memory::new(self.init_memory()?);

    Instance::validate(&func.params, &args)?;
    let mut runtime = Runtime::new(function_table, memory);
    let mut stack = runtime.execute(index, &args)?;

    stack.reverse();
    let ret = &stack[0..return_length];
    Ok(ret.to_vec())
}

execute

関数名から実行すべきコードの解決や引数のvalidationが出来たので実行していくだけになりました。
インストラクションポインタをインクリメントしつつ命令を実行していきます。ここからは実装を頑張るだけになってきますね。

余談ですが、ifや関数呼び出しが合ったときにはインストラクションポインタをガチャガチャ更新していたのですがバグを多く出して悲しい気持ちになりました。まだまだバグが見つかりそうで怖いです。

pub fn execute(
    &mut self,
    func_index: usize,
    args: &[RuntimeValue],
) -> Result<ValueStack, RuntimeError> {
    // 関数呼び出し時にローカル変数のスコープをうまいこと切りたいのでactivation stackというものを導入している。
    // ローカル変数はこのstackに属するようになっているのでスコープ外の変数は見えない
    self.activation_stack = ActivationStack::init(func_index, args.to_vec());

    while let Some(instruction) = self.get_instruction()? {
        self.increment_pc()?;

        match instruction {
            Instruction::Reserved => {}
            Instruction::Prefix(_) => {}
            Instruction::Nop => {}
            Instruction::I32Add => self.add::<i32>(),
            Instruction::I32Sub => self.sub::<i32>(),
            Instruction::I32Mul => self.mul::<i32>(),
            ... // 以下大量のmatch
        }
    }
    Ok(self.value_stack.clone())
}

まとめ + 感想

だいぶ駆け足になってしまったのと、decoderもexecutorも概要だけの説明になってしまいました、、、、

wasmインタプリタは3000行ちょっとでいい感じのところまで動いてくれるので趣味プロジェクトとしていい感じのサイズ感でした。
またテストツールも充実しているので開発体験が良かったです。

:) % tokei ./src
===============================================================================
 Language            Files        Lines         Code     Comments       Blanks
===============================================================================
 Rust                   18         3135         2738           22          375
 |- Markdown             2            2            0            2            0
 (Total)                           3137         2738           24          375
===============================================================================
 Total                  18         3135         2738           22          375
===============================================================================

このプロジェクトの当初の目的はwasmを完全に理解してwasmのプロポーザルを呼んだり、wasmtimeというruntimeを読むことだったのですが、完全に力尽きてしまいました、、、が今回がインタプリタを作るのは初めてだったので結構良い経験になりました。

WebAssemblyが色んな所に使われていく未来を想像すると熱い気持ちになるのでこれからもwasmの動向を見ていきたいなーと思っている今日このごろです。
それではおやすみなさい。