paritytech / wasm-instrument

Instrument and transform wasm modules.
Apache License 2.0
46 stars 15 forks source link

get rid of confusing `ACTIVATION_FRAME_COST` constant #80

Open StackOverflowExcept1on opened 10 months ago

StackOverflowExcept1on commented 10 months ago

Hi, I'm working on adding a custom stack limiter for our protocol. It will provide API like that: (e.g. you can inject something instead of unreachable instruction)

wasm_instrument::inject_custom_stack_limiter(module, 20_000, |signature| vec![
    Instruction::Foo,
    Instruction::Bar,
    ...,
])

And I think I understood what the real problem was at #1

https://github.com/paritytech/wasm-instrument/blob/6307588b3dc508d8b24ee64c364841d89e558ec2/src/stack_limiter/mod.rs#L10-L19

So we have an instrumented_call! with a maximum stack height of 2, the same as the ACTIVATION_FRAME_COST constant.

// stack_height += stack_cost(F) 
GetGlobal($stack_height_global_idx), // push #1
I32Const($callee_stack_cost), // push #2
// max_height = 2

Backing to problem #1: the function with index 0 has no arguments and no return value, so its max_height=0

(module
  (type (;0;) (func))
  (func (;0;) (type 0)
    call 0)
  (func (;1;) (type 0)
    call 0)
  (export "main" (func 1)))

We also have a main export that will call function 0 and it will not be instrumented due to the condition: https://github.com/paritytech/wasm-instrument/blob/6307588b3dc508d8b24ee64c364841d89e558ec2/src/stack_limiter/thunk.rs#L44-L48

What is the best way to solve this problem? When the Call(idx) opcode is found, we need to check that idx is not an import and simply do:

let import_count = module.import_count(elements::ImportCountType::Function);
let mut max_height: u32 = 0;
let mut overhead_of_instrumented_call = 0;

...

match opcode {
    Call(idx) => {
        if idx >= import_count {
            overhead_of_instrumented_call = 2;
        }
    }
}

Ok(max_height + overhead_of_instrumented_call)

or we can start passing instrument_call! instructions instead of Call(idx) to match opcode

athei commented 10 months ago

But even with non-imports the problem can arise if they have a zero cost without the constant?

Also, you still need the constant if the callee is an import?

StackOverflowExcept1on commented 10 months ago

But even with non-imports the problem can arise if they have a zero cost without the constant?

It should not because we need to charge for overhead of the instrumented call instructions. For example, instead of calculating the stack height for Call(0) we will replace the Call(0) with info[0] (sequence that we got from instrumented call! macro)

In this case the Call(idx), which does not change the stack (0 pops and 0 pushes), will still cost 2.

#![feature(if_let_guard)]

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Instruction {
    Foo,
    Bar,
    Baz,
    Call(usize),
}

fn process(opcode: &Instruction) {
    println!("{opcode:?}");
}

fn main() {
    use Instruction::*;

    let info = vec![vec![Foo, Bar, Baz]];
    let xs = vec![Foo, Bar, Baz, Call(0), Foo];

    for opcode in xs.iter() {
        match opcode {
            Call(i) if let Some(slice) = info.get(*i) => {
                for opcode in slice {
                    process(opcode);
                }
            }
            _ => process(opcode),
        }
    }
}
StackOverflowExcept1on commented 10 months ago

Also, you still need the constant if the callee is an import?

Why do I need this constant if instrument_function will not process imports (since they have zero cost)?

https://github.com/paritytech/wasm-instrument/blob/0c739d92c453c2a1ff8a2a3c44e7299ff446ea94/src/stack_limiter/mod.rs#L164-L165

https://github.com/paritytech/wasm-instrument/blob/0c739d92c453c2a1ff8a2a3c44e7299ff446ea94/src/stack_limiter/mod.rs#L260-L266

For example, we have module

(module
  (import "env" "nop" (func $nop))
  (func $f1 (export "f1")
    call $nop
  )
)

After instrumentation:

(module
  (type (;0;) (func))
  (import "env" "nop" (func $nop (;0;) (type 0))) ;; cost of import = 0
  (func $f1 (;1;) (type 0) ;; cost of func = 0
    call $nop ;; it will not be processed by the `instrument_function`
  )
  (global (;0;) (mut i32) i32.const 0)
  (export "f1" (func $f1)) ;; thunk function is not created here, bcz we can safely call to original
)
athei commented 10 months ago

What I am saying is that this is not only about imports is it? You can just call an zero cost function recursively without a constant cost for the activation.

StackOverflowExcept1on commented 10 months ago

I think you mean that example with empty functions, right? If you look at the PR, the tests pass and I'll explain why.

So we have function 0, which actually costs 0 (no params, nor results)

(func (;0;)
  call 0
)

but after the injection stack limiter it will look like this:

(func (;0;) (type 0)
  global.get 0
  i32.const 2
  i32.add
  global.set 0
  global.get 0
  i32.const 1024
  i32.gt_u
  if ;; label = @1
    unreachable
  end
  call 0
  global.get 0
  i32.const 2
  i32.sub
  global.set 0
)

Instead of counting the cost of Call(idx=0) we will count cost of instrument_call!(idx, 0, 0, 0). Therefore this function will have a maximum stack height of 2, not 0.