Export the rust version to WASM + Use it in a Go library

Author: ngotchac

Makefile

@@ -0,0 +1,5 @@
+.PHONY: wasm
+wasm:
+	cd rust && \
+    cargo build --target wasm32-unknown-unknown --release
+	cp rust/target/wasm32-unknown-unknown/release/evmole.wasm go/evmole.wasm

evmole.wasm

@@ -0,0 +1,75 @@
+package evmole
+
+import (
+	"context"
+	_ "embed"
+	"encoding/binary"
+	"fmt"
+
+	"github.com/tetratelabs/wazero"
+)
+
+//go:embed evmole.wasm
+var evmoleWASM []byte
+
+// FunctionSelectors tries to extract from the given bytecode the function selectors.
+func FunctionSelectors(ctx context.Context, code []byte) ([][4]byte, error) {
+	// Create a new WebAssembly runtime
+	runtime := wazero.NewRuntime(ctx)
+	defer func() { _ = runtime.Close(ctx) }()
+
+	instance, err := runtime.Instantiate(ctx, evmoleWASM)
+	if err != nil {
+		panic(fmt.Errorf("failed to instantiate WASM module: %w", err))
+	}
+	defer func() { _ = instance.Close(ctx) }()
+
+	gasLimit := 0
+
+	memory := instance.Memory()
+	functionSelectorsFunc := instance.ExportedFunction("ext_function_selectors")
+
+	codeOffset := uint32(0)
+	resultLenOffset := uint32(len(code))
+	resultOffset := resultLenOffset + 4
+	resultCapacity := uint32(512 * 4)
+
+	// Write input to memory
+	ok := memory.Write(codeOffset, code)
+	if !ok {
+		return nil, fmt.Errorf("failed to write input to memory")
+	}
+
+	// Call the WASM function
+	results, err := functionSelectorsFunc.Call(ctx,
+		uint64(codeOffset), uint64(len(code)), uint64(gasLimit),
+		uint64(resultLenOffset), uint64(resultOffset), uint64(resultCapacity))
+	if err != nil {
+		return nil, fmt.Errorf("failed to call function_selectors: %w", err)
+	}
+
+	if status := uint32(results[0]); status != 0 {
+		return nil, fmt.Errorf("error: status=%d", status)
+	}
+
+	// Read the actual result length
+	rawResultLen, ok := memory.Read(resultLenOffset, 4)
+	if !ok {
+		return nil, fmt.Errorf("failed to read result length")
+	}
+	resultLen := binary.LittleEndian.Uint32(rawResultLen)
+
+	// Success, read the result
+	result, ok := memory.Read(resultOffset, resultLen)
+	if !ok {
+		return nil, fmt.Errorf("failed to read result from memory")
+	}
+
+	// Convert to [][4]byte
+	selectors := make([][4]byte, len(result)/4)
+	for i := 0; i < len(result); i += 4 {
+		copy(selectors[i/4][:], result[i:i+4])
+	}
+
+	return selectors, nil
+}

go/evmole.go

@@ -0,0 +1,5 @@
+module github.com/cdump/evmole/go
+
+go 1.22.5
+
+require github.com/tetratelabs/wazero v1.7.3

go/evmole.wasm

@@ -0,0 +1,2 @@
+github.com/tetratelabs/wazero v1.7.3 h1:PBH5KVahrt3S2AHgEjKu4u+LlDbbk+nsGE3KLucy6Rw=
+github.com/tetratelabs/wazero v1.7.3/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y=

go/evmole_test.go

@@ -8,6 +8,9 @@ license = "MIT"
 readme = "README.md"
 repository = "https://github.com/cdump/evmole"
 
+[lib]
+crate-type = ["cdylib"]
+
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]

go/go.mod

@@ -14,3 +14,34 @@ mod evm;
 pub mod selectors;
 
 pub type Selector = [u8; 4];
+
+#[no_mangle]
+pub extern "C" fn ext_function_selectors(
+    code_ptr: *const u8,
+    code_len: usize,
+    gas_limit: u32,
+    result_len_ptr: *mut u32,
+    result_ptr: *mut u8,
+    result_capacity: usize,
+) -> u32 {
+    let code = unsafe { std::slice::from_raw_parts(code_ptr, code_len) };
+    let selectors = function_selectors(code, gas_limit);
+
+    let flattened: Vec<u8> = selectors.into_iter().flatten().collect();
+    let result_len = flattened.len();
+
+    unsafe {
+        *result_len_ptr = result_len as u32;
+    }
+
+    if result_len > result_capacity {
+        // Buffer too small
+        return 1;
+    }
+
+    unsafe {
+        std::ptr::copy_nonoverlapping(flattened.as_ptr(), result_ptr, result_len);
+    }
+    // Success
+    0
+}