wasabi icon indicating copy to clipboard operation
wasabi copied to clipboard

resolveTableIdx gives wrong function index when using multiple wasm modules

Open doehyunbaek opened this issue 1 year ago • 5 comments

Repository reproducing the problem: https://github.com/doehyunbaek/wasabi_multimodule_fidx

Currently, resolveTableIdx relies on name property of the function to resolve function index.

https://github.com/danleh/wasabi/blob/21a322b7faac9440b931762aae124ffa57d0fa17/crates/wasabi/js/runtime.js#L45-L51

This does not work when the function is imported from other WebAssembly instance and the index of the function in its imported module and the index of the function in its declared module is different, e.g.

module1.wat

(module
    (func $private_func_0 (;0;) )
    (func $private_func_1 (;1;) )
    (func $shared_func (;2;) (param $0 i32) (result i32)
        local.get $0
        i32.const 1
        i32.add
    )
    (export "shared_func" (func $shared_func))
)

module2.wat

(module
    (import "env" "shared_func" (func $shared_func (;0;) ))
)

For module2, resolveTableIdx would resolve shared_func to 2, which is different from expected behavior, which is 0.

doehyunbaek avatar Jul 28 '24 10:07 doehyunbaek

Thanks, great analysis and very concise reproducer! :) This is indeed a bug.

Besides the multi-module case you layed out here, I'm wondering what will happen for functions that were created with the WebAssembly.Function API (still a proposal, but already implemented in some engines/browsers). I bet it's wrong there too.

How do we solve this? Basically, how to map back from a table entry to a function index in the local module? Even more problematic: Can we even always map back? E.g., what happens if we have a Wasm module that exports its table and the host inserts (via WebAssembly.Table.set) a function that is not even declared in the module (neither as import nor with code)? I guess this is possible host behavior, and in that case we cannot return a meaningful index, just null.

So the proper solution is probably:

  • Expose the function reference directly to the hook (i.e., add a func_ref argument to after indirectTableIdx in the call_pre hook), and return targetFunc == undefined if we could not map the reference back to a local function index.
  • For those cases, where the call target is a function from another module, but is also properly imported into the local module, we should somehow return the local index. I'm not quite sure what to do, brainstorming ideas: a) Compare the function reference against all functions in the module (which we could export as part of the instrumentation). Obvious downside: each call would do O(n) comparisons :( b) Maybe we can have a fastpath: We keep a map from indices to function references, look up the reference from the index of the Function object, and if that is the same object, we return the index, otherwise we do the slowpath? I.e., something like
    const func_idx_to_ref = { 0: instance.exports._wasabi_func_0, 1: instance.exports._wasabi_func_1, ... }
    // some helper to get the function index in the call_pre hook:
    let func_ref = table.get(table_index);
    let maybe_correct_func_idx = parseInt(func_ref.name);
    if (func_ref == func_idx_to_ref(maybe_correct_func_idx)) {
      // local function, works:
      return maybe_correct_func_idx;
    } else {
      for (let [idx, ref] of Object.entries(func_idx_to_ref)) {
        if (ref == func_ref) return idx;
      }
      return undefined;
    }
    
    c) Somehow track the state of the table with a "shadow table" that contains function indices instead of references. Downside: breaks if the host mutates the table (or with Wasm extension that adds table.get/table.set). On the other hand, host code that mutates the table is also wrong today already...

WDYT?

danleh avatar Jul 28 '24 17:07 danleh

Mhm, maybe even simpler: one can just have a Map from function references to index, and either prefill at initialization time with all references in the table or lazily initialize it with the "slowpath" code from above.

danleh avatar Jul 28 '24 17:07 danleh

b) Maybe we can have a fastpath

I like this approach best, will tinker around with this in mind.

Map from function references to index

So it's a funcref --> fidx instead of fidx --> funcref as above? I think this is better, although I'm not sure we can construct such a map having funcref as a key. Will try and update.

doehyunbaek avatar Jul 29 '24 01:07 doehyunbaek

Map from function references to index

This works! I played around with the https://github.com/sola-st/wasm-r3/commit/1004e657bb83f1c5170bc60464ffe7a19fb30d71 commit and it seems it works well at least for the Wasm-R3's use case.

I want to generalize this solution and upstream to Wasabi but I think it is better to merge any outstanding changes(esp. https://github.com/danleh/wasabi/pull/41) and build on top of that, WDYT?

doehyunbaek avatar Jul 29 '24 04:07 doehyunbaek

Yes, sounds good (merging first, unless fixing this is urgent)

danleh avatar Jul 29 '24 08:07 danleh