ZEVORN.blog

October 8, 2024

基于 Rust 编写最小 JIT

article2.6 min to read

即时编译（Just-In-Time Compilation，简称 JIT 编译）是一种动态编译技术，在程序运行期间将高级语言的源代码或者中间表示形式（如字节码）转换为机器代码。与传统的静态编译不同，JIT 编译器在程序运行时选择性地编译代码片段，并且可以根据运行时的具体情况优化这些代码。

下面使用 Rust 构建一个最小 JIT 案例，参考 rustyjit 。

基本思路

  1. 在程序进程运行时 mmap 一块内存 code_buffer（4KiB 对齐），权限设置为可读可写可执行；
  2. 将宿主机指令序列的二进制编码，发射到 code_buffer；
  3. 将 coder_buffer 中的程序段入口地址，强转为一个函数指针；
  4. 执行这个函数指针。

前置知识

FFI 编程

在进程中申请内存，需要调用 libc 库的接口，这涉及到 Rust FFI 编程 - libc crate。简单讲，我们需要 extern Rust 封装好的 libc 的 crate，从而获取我们需要的接口。主要如下：

  1. mmap()
  2. memset()
  3. mprotect()
  4. memalign()

注意有一些接口，在 Rust 中是封装好的变种方法，一般是关联函数，不能使用原生 API 名称。

特征 Trait

特征 Trait 可以告诉编译器一个特定的类型所具有的、且能跟其它类型共享的特性。我们可以使用特征通过抽象的方式来定义这种共享行为，还可以使用特征约束来限定一个泛型类型必须要具有某个特定的行为。

由于 Rust 是一个安全性极高的语言，我们需要依靠 Trait，来实现内存的可变访问。

构建最小 JIT

首先定义一个结构体，它包含一个指向 code_buffer 的指针：

struct JitMemory {    code_buffer: *mut u8,}

然后我们编写一个 new 方法，实现内存申请与初始化：

std::mem;// 定义页面大小常量const PAGE_SIZE: usize = 4096;impl JitMemory {    fn new(num_pages: usize) -> JitMemory {        let code_buffer: *mut u8;        unsafe {            let size: usize = num_pages * PAGE_SIZE;            let mut _contents: *mut libc::c_void =            mem::MaybeUninit::<libc::c_void>::uninit().as_mut_ptr();            // 分配内存并对齐到页面大小            libc::posix_memalign(&mut _contents, PAGE_SIZE, size);            // 设置内存权限为可执行、可读、可写            libc::mprotect(                _contents,                size,                libc::PROT_EXEC | libc::PROT_READ | libc::PROT_WRITE,            );            // 初始化内存内容为 'RET' 指令            libc::memset(_contents, 0xc3, size);            code_buffer = mem::transmute(_contents);        }        JitMemory { code_buffer }    }}

接着我们实现指令二进制编码数据的字节码发射功能：

use std::ops::{Index, IndexMut};// 实现 Index trait，允许通过索引访问内存impl Index<usize> for JitMemory {    type Output = u8; // 声明了 Index<usize> trait 的关联类型 Output 为 u8 类型    fn index(&self, _index: usize) -> &u8 {        unsafe { &*self.code_buffer.offset(_index as isize) }    }}// 实现 IndexMut trait，允许通过索引修改内存impl IndexMut<usize> for JitMemory {    fn index_mut(&mut self, _index: usize) -> &mut u8 {        unsafe { &mut *self.code_buffer.offset(_index as isize) }    }}

最后我们编写一段简单的汇编程序，将其发射到内存并运行：

fn run_jit() -> fn() -> i64 {    let mut jit: JitMemory = JitMemory::new(1);    // 在 JIT 内存中写入机器码，生成一个返回 3 的函数    jit[0] = 0x48; // mov RAX, 0x30    jit[1] = 0xc7;    jit[2] = 0xc0;    jit[3] = 0x30;    jit[4] = 0;    jit[5] = 0;    jit[6] = 0;    // 将内存指针转换为函数指针并返回    unsafe { mem::transmute(jit.code_buffer) }}fn main() {    // 生成并调用 JIT 编译的函数，并打印结果    let fun: fn() -> i64 = run_jit();    println!("{:#x}", fun());}

编译运行，看看输出结果：

$ cargo run...0x30