75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
import ctypes
|
|
import struct
|
|
|
|
import wasmtime
|
|
|
|
from phasmplatform.common.router import BaseRouter
|
|
from .base import BaseRunner
|
|
|
|
|
|
class WasmTimeRunner(BaseRunner):
|
|
__slots__ = ('store', 'module', 'instance', 'exports')
|
|
|
|
def __init__(self, router: BaseRouter, wasm_bin: bytes) -> None:
|
|
super().__init__(router)
|
|
|
|
self.store = wasmtime.Store()
|
|
|
|
post_message_bind = wasmtime.Func(self.store, wasmtime.FuncType([
|
|
wasmtime.ValType.i32(),
|
|
wasmtime.ValType.i32(),
|
|
wasmtime.ValType.i32(),
|
|
wasmtime.ValType.i32(),
|
|
], []), self.post_message)
|
|
|
|
self.module = wasmtime.Module(self.store.engine, wasm_bin)
|
|
self.instance = wasmtime.Instance(self.store, self.module, [
|
|
post_message_bind,
|
|
])
|
|
|
|
self.exports = self.instance.exports(self.store)
|
|
|
|
def alloc_bytes(self, data: bytes) -> int:
|
|
memory = self.exports['memory']
|
|
assert isinstance(memory, wasmtime.Memory) # type hint
|
|
|
|
data_ptr = memory.data_ptr(self.store)
|
|
data_len = memory.data_len(self.store)
|
|
|
|
alloc_bytes = self.exports['stdlib.types.__alloc_bytes__']
|
|
assert isinstance(alloc_bytes, wasmtime.Func)
|
|
|
|
ptr = alloc_bytes(self.store, len(data))
|
|
assert isinstance(ptr, int) # type hint
|
|
|
|
idx = ptr + 4 # Skip the header from header from __alloc_bytes__
|
|
for byt in data:
|
|
assert idx < data_len
|
|
data_ptr[idx] = ctypes.c_ubyte(byt)
|
|
idx += 1
|
|
|
|
return ptr
|
|
|
|
def read_bytes(self, ptr: int) -> bytes:
|
|
memory = self.exports['memory']
|
|
assert isinstance(memory, wasmtime.Memory) # type hint
|
|
|
|
data_ptr = memory.data_ptr(self.store)
|
|
data_len = memory.data_len(self.store)
|
|
|
|
raw = ctypes.string_at(data_ptr, data_len)
|
|
|
|
length, = struct.unpack('<I', raw[ptr:ptr + 4]) # Header prefixed by __alloc_bytes__
|
|
|
|
return raw[ptr + 4:ptr + 4 + length]
|
|
|
|
def handle_message(self, namespace: bytes, topic: bytes, kind: bytes, body: bytes) -> None:
|
|
namespace_ptr = self.alloc_bytes(namespace)
|
|
topic_ptr = self.alloc_bytes(topic)
|
|
kind_ptr = self.alloc_bytes(kind)
|
|
body_ptr = self.alloc_bytes(body)
|
|
|
|
handle_message = self.exports['handle_message']
|
|
assert isinstance(handle_message, wasmtime.Func)
|
|
handle_message(self.store, namespace_ptr, topic_ptr, kind_ptr, body_ptr)
|