wasmtime / issue #4466 Out-of-band fuel metering · git-wasmtime

struct VMRuntimeLimits {
    // ...

    /// A flag that indicates if the out-of-band fuel metering is enabled.
    pub outband_fuel: UnsafeCell<bool>,

    /// host-to-wasm return address.
    pub orig_entry_trampoline_return_address: UnsafeCell<usize>,

    /// wasm-to-host return address.
    pub orig_exit_trampoline_return_address: UnsafeCell<usize>,
}

// ...

asm_func!(
    "host_to_wasm_trampoline",
    "
        .cfi_startproc simple
        .cfi_def_cfa_offset 0

        // Load the pointer to `VMRuntimeLimits` in `r10`.
        mov r10, 8[rsi]

        // Check to see if this is a core `VMContext` (MAGIC == 'core').
        cmp DWORD PTR [rdi], 0x65726f63

        // Store the last Wasm SP into the `last_wasm_entry_sp` in the limits, if this
        // was core Wasm, otherwise store an invalid sentinal value.
        mov r11, -1
        cmove r11, rsp
        mov 40[r10], r11

        // Check if the out-of-band fuel metering is enabled. If it is not enabled, skip to the
        // tail call directly.
        cmp QWORD PTR 48[r10], 0x00
        je host_to_wasm_trampoline__perform_tail_call

        // OK, here we know that the out-of-band fuel metering is enabled.
        //
        // Dump the `fuel_consumed` into the fuel register.
        mov r15, 8[r10]

        // Next, we have to do some shenanigans. Below we call into the wasm with a tail call, to
        // not mess with the stack layout. But at the same time, after returning from wasm, we
        // have to our fuel register r15 back into `fuel_consumed` of `VMRuntimeLimits`. To do that
        // we stash the original return address (which is currently located at `[rsp]`) that points
        // back to the caller from the host side and change it with the continuation defined
        // below. The continuation will dump the r15 back into the memory and unstash the original
        // return address and jump at it.

        mov r11, [rsp]
        mov 56[r10], r11
        lea r11, [rip+host_to_wasm_trampoline__cont]
        mov [rsp], r11

        host_to_wasm_trampoline__perform_tail_call:

        // Tail call to the callee function pointer in the vmctx.
        jmp 16[rsi]

        host_to_wasm_trampoline__cont:

        // Here we're in a inconvenient situation, where the callee's frame is popped and all the
        // arguments are gone. We have to improvise our way out of this situation.

        // Open a temporary stack frame.
        //
        // The frame layout:
        //
        // rsp +  0: The original return address. Written by `wasmtime_outband_fuel_on_entry`.
        // rsp +  8: r15
        // rsp + 16: rax
        // rsp + 24: padding
        // rsp + 32: xmm0
        sub rsp, 48
        mov 8[rsp], r15
        mov 16[rsp], rax
        movdqu XMMWORD PTR 32[rsp], xmm0

        // Finally, move the current rsp into the first argument.
        mov rdi, rsp
        call wasmtime_outband_fuel_on_entry

        // Load the original return address into r10
        mov r10, 0[rsp]

        // Restore the initial values of rax and xmm0.
        mov rax, 16[rsp]
        movdqu xmm0, XMMWORD PTR 32[rsp]

        // Release the allocated stack area.
        add rsp, 48

        // Jump at the original return address.
        jmp r10

        .cfi_endproc
    ",
);


/// Dumps fuel into the `VMRuntimeLimits->fuel_consumed` and unstashes the original return address
/// for the trampoline.
///
/// Only used if out-of-band fuel is used.
#[no_mangle]
pub unsafe extern "C" fn wasmtime_outband_fuel_on_entry(sp: *mut usize) {
    use crate::traphandlers::tls;
    tls::with(|info| {
        let info = match info {
            Some(info) => info,
            None => {
                // Entering trampoline must set the TLS.
                std::hint::unreachable_unchecked();
            }
        };

        #[repr(C)]
        struct StackArgs {
            return_address: usize,
            fuel: i64,
        }

        let args = sp as *mut StackArgs;

        (*args).return_address = *(*info.runtime_limits())
            .orig_entry_trampoline_return_address
            .get();
        *(*info.runtime_limits()).fuel_consumed.get() = (*args).fuel;

        if (*args).fuel > 0 {
            // Fuel ran out and we bail through unwinding.
            crate::traphandlers::raise_lib_trap(wasmtime_environ::TrapCode::Interrupt);
        }
    });
}

// extern crate backtrace;

use std::arch;
use std::mem;

arch::global_asm!(
    "
    .p2align 4
    .global trampoline
    .type trampoline, @function

    trampoline:

    .cfi_startproc simple
    .cfi_def_cfa rsp, 0

    // move RDI, the original function, into r11
    mov r11, rdi

    // the return address pushed by the `call` instruction calling this trampoline
    // is located at *rsp. Move the value to rdi.
    mov rdi, [rsp]

    // Then, replace the return value with the continuation at the label `cont`.
    lea r10, cont[rip]
    mov [rsp], r10

    //
    // Important!!
    //
    // Mark the register 16, aka the return address, as undefined. That stops the unwinder here.
    // If that was not here, the unwinder would fall into an infinite loop.
    //
    // An unfortunate side-effect is the stack traces are also chopped here.
    .cfi_undefined 16

    // Finally, perform the jump.
    jmp r11

    cont:

    jmp rax

    .cfi_endproc

    .size trampoline, .-trampoline
    "
);

extern "C" {
    fn trampoline();
}

type Func = extern "C" fn(*const (), usize) -> *const ();

#[inline]
unsafe fn prepare_trampoline(dest: Func) -> (*const (), Func) {
    let dest = dest as *const Func as *const ();
    let func = mem::transmute_copy(&(trampoline as usize));

    (dest, func)
}

fn main() {
    unsafe {
        let (dest, f) = prepare_trampoline(host_func);
        f(dest, 48);
    }
}

extern "C" fn host_func(rdi: *const (), rsi: usize) -> *const () {
    println!("{}", rsi);
    let mut count = 0;
    backtrace::trace(|frame| {
        let ip = frame.ip();
        println!("{:X}", ip as usize);

        if count == 100 {
            // likely we stuck in a loop. Break into the debugger.
            unsafe {
                arch::asm!("ud2");
            }
        }
        count += 1;

        true
    });

    // we don't have such luxury in the real thing.
    rdi
}

* thread #1, name = 'cfi-dwarf', stop reason = signal SIGABRT
  * frame #0: 0x00007ffff7e3abaa libc.so.6`raise + 202
    frame #1: 0x00007ffff7e25523 libc.so.6`abort + 255
    frame #2: 0x00007ffff7dd5b40 libgcc_s.so.1`uw_update_context_1 + 1008
    frame #3: 0x00007ffff7dd5cd1 libgcc_s.so.1`uw_update_context + 17
    frame #4: 0x00007ffff7dd68bd libgcc_s.so.1`_Unwind_Backtrace + 93
    frame #5: 0x000055555555cfa1 cfi-dwarf`backtrace::backtrace::trace_unsynchronized::hf1601be64f8e51c8 [inlined] backtrace::backtrace::libunwind::trace::h80afc20a98e10b83(cb=&mut dyn core::ops::function::FnMut<(&backtrace::backtrace::Frame), Output=bool> @ 0x00007fffffff9d58) at libunwind.rs:93:5
    frame #6: 0x000055555555cf8c cfi-dwarf`backtrace::backtrace::trace_unsynchronized::hf1601be64f8e51c8(cb={closure_env#0} @ 0x00007fffffff9d40) at mod.rs:66:5
    frame #7: 0x000055555555d038 cfi-dwarf`backtrace::backtrace::trace::h4ba2f0e531a2363e(cb={closure_env#0} @ 0x00007fffffff9da0) at mod.rs:53:14
    frame #8: 0x000055555555d1d3 cfi-dwarf`cfi_dwarf::host_func::h89c8978f2a6c15a6(rdi=0x000055555555d148, rsi=48) at main.rs:74:5

Stream: git-wasmtime

Topic: wasmtime / issue #4466 Out-of-band fuel metering

Wasmtime GitHub notifications bot (Jul 19 2022 at 18:00):

Subscribe to Label Action

Wasmtime GitHub notifications bot (Jul 19 2022 at 18:01):

Label Messager: wasmtime:config

Example

Wasmtime GitHub notifications bot (Jul 22 2022 at 16:26):

Wasmtime GitHub notifications bot (Jul 22 2022 at 19:30):

Wasmtime GitHub notifications bot (Jul 25 2022 at 15:14):

Wasmtime GitHub notifications bot (Jul 25 2022 at 15:36):

Wasmtime GitHub notifications bot (Jul 25 2022 at 15:36):

Wasmtime GitHub notifications bot (Jul 25 2022 at 16:02):

Wasmtime GitHub notifications bot (Jul 25 2022 at 16:02):

Wasmtime GitHub notifications bot (Jul 25 2022 at 21:43):

Wasmtime GitHub notifications bot (Jul 26 2022 at 09:58):

Wasmtime GitHub notifications bot (Jul 26 2022 at 16:14):

Wasmtime GitHub notifications bot (Jul 26 2022 at 16:41):

Wasmtime GitHub notifications bot (Jul 26 2022 at 19:59):