wazero/internal/wasm/host.go

package internalwasm

import (
	"context"
	"encoding/binary"
	"fmt"
	"math"

	publicwasm "github.com/tetratelabs/wazero/wasm"
)

// compile time check to ensure ModuleContext implements publicwasm.ModuleContext
var _ publicwasm.ModuleContext = &ModuleContext{}

func NewModuleContext(s *Store, instance *ModuleInstance) *ModuleContext {
	return &ModuleContext{
		Engine: s.Engine,
		ctx:    context.Background(),
		memory: instance.Memory,
		Module: instance,
	}
}

// ModuleContext implements wasm.ModuleContext and wasm.Module
type ModuleContext struct {
	// Engine is exported for wazero.MakeWasmFunc
	Engine Engine
	// Module is exported for wazero.MakeWasmFunc
	Module *ModuleInstance
	// memory is exposed as wasm.ModuleContext Memory
	memory publicwasm.Memory
	// ctx is exposed as wasm.ModuleContext Context
	ctx context.Context
}

// WithContext allows overriding context without re-allocation when the result would be the same.
func (c *ModuleContext) WithContext(ctx context.Context) *ModuleContext {
	// only re-allocate if it will change the effective context
	if ctx != nil && ctx != c.ctx {
		return &ModuleContext{Engine: c.Engine, Module: c.Module, memory: c.memory, ctx: ctx}
	}
	return c
}

// WithMemory allows overriding memory without re-allocation when the result would be the same.
func (c *ModuleContext) WithMemory(memory *MemoryInstance) *ModuleContext {
	// only re-allocate if it will change the effective memory
	if memory != nil && memory.Max != nil && *memory.Max > 0 && memory != c.memory {
		return &ModuleContext{Engine: c.Engine, Module: c.Module, memory: memory, ctx: c.ctx}
	}
	return c
}

func (c *ModuleContext) Context() context.Context {
	return c.ctx
}

// Memory implements wasm.Module Memory
func (c *ModuleContext) Memory() publicwasm.Memory {
	return c.memory
}

// Function implements wasm.Functions Function
func (c *ModuleContext) Function(name string) (publicwasm.Function, bool) {
	exp, err := c.Module.GetExport(name, ExportKindFunc)
	if err != nil {
		return nil, false
	}
	return (&function{c: c, f: exp.Function}).Call, true
}

type function struct {
	c *ModuleContext
	f *FunctionInstance
}

func (f *function) Call(ctx context.Context, params ...uint64) ([]uint64, error) {
	paramSignature := f.f.FunctionType.Type.Params
	paramCount := len(params)
	if len(paramSignature) != paramCount {
		return nil, fmt.Errorf("expected %d params, but passed %d", len(paramSignature), paramCount)
	}
	hc := f.c.WithContext(ctx)
	return hc.Engine.Call(hc, f.f, params...)
}

// ExportHostFunctions is defined internally for use in WASI tests and to keep the code size in the root directory small.
func (s *Store) ExportHostFunctions(moduleName string, nameToGoFunc map[string]interface{}) (publicwasm.HostExports, error) {
	if err := s.requireModuleUnused(moduleName); err != nil {
		return nil, err
	}

	ret := HostExports{NameToFunctionInstance: make(map[string]*FunctionInstance, len(nameToGoFunc))}
	for name, goFunc := range nameToGoFunc {
		if hf, err := NewGoFunc(name, goFunc); err != nil {
			return nil, err
		} else if function, err := s.AddHostFunction(moduleName, hf); err != nil {
			return nil, err
		} else {
			ret.NameToFunctionInstance[name] = function
		}
	}
	return &ret, nil
}

func (s *Store) requireModuleUnused(moduleName string) error {
	if _, ok := s.hostExports[moduleName]; ok {
		return fmt.Errorf("module %s has already been exported by this host", moduleName)
	}
	if _, ok := s.ModuleContexts[moduleName]; ok {
		return fmt.Errorf("module %s has already been instantiated", moduleName)
	}
	return nil
}

// HostExports implements wasm.HostExports
type HostExports struct {
	NameToFunctionInstance map[string]*FunctionInstance
}

// call implements wasm.HostFunction
func (f *FunctionInstance) call(ctx publicwasm.ModuleContext, params ...uint64) ([]uint64, error) {
	hCtx, ok := ctx.(*ModuleContext)
	if !ok { // TODO: guard that hCtx.Module actually imported this!
		return nil, fmt.Errorf("this function was not imported by %s", ctx)
	}
	return hCtx.Engine.Call(hCtx, f, params...)
}

// Function implements wasm.HostExports Function
func (g *HostExports) Function(name string) (publicwasm.HostFunction, bool) {
	f, ok := g.NameToFunctionInstance[name]
	return f.call, ok
}

// Len implements wasm.ModuleContext Len
func (m *MemoryInstance) Len() uint32 {
	return uint32(len(m.Buffer))
}

// hasLen returns true if Len is sufficient for sizeInBytes at the given offset.
func (m *MemoryInstance) hasLen(offset uint32, sizeInBytes uint32) bool {
	return uint64(offset+sizeInBytes) <= uint64(m.Len()) // uint64 prevents overflow on add
}

// ReadUint32Le implements wasm.ModuleContext ReadUint32Le
func (m *MemoryInstance) ReadUint32Le(offset uint32) (uint32, bool) {
	if !m.hasLen(offset, 4) {
		return 0, false
	}
	return binary.LittleEndian.Uint32(m.Buffer[offset : offset+4]), true
}

// ReadFloat32Le implements wasm.ModuleContext ReadFloat32Le
func (m *MemoryInstance) ReadFloat32Le(offset uint32) (float32, bool) {
	v, ok := m.ReadUint32Le(offset)
	if !ok {
		return 0, false
	}
	return math.Float32frombits(v), true
}

// ReadUint64Le implements wasm.ModuleContext ReadUint64Le
func (m *MemoryInstance) ReadUint64Le(offset uint32) (uint64, bool) {
	if !m.hasLen(offset, 8) {
		return 0, false
	}
	return binary.LittleEndian.Uint64(m.Buffer[offset : offset+8]), true
}

// ReadFloat64Le implements wasm.ModuleContext ReadFloat64Le
func (m *MemoryInstance) ReadFloat64Le(offset uint32) (float64, bool) {
	v, ok := m.ReadUint64Le(offset)
	if !ok {
		return 0, false
	}
	return math.Float64frombits(v), true
}

// Read implements wasm.ModuleContext Read
func (m *MemoryInstance) Read(offset, byteCount uint32) ([]byte, bool) {
	if !m.hasLen(offset, byteCount) {
		return nil, false
	}
	return m.Buffer[offset : offset+byteCount], true
}

// WriteUint32Le implements wasm.ModuleContext WriteUint32Le
func (m *MemoryInstance) WriteUint32Le(offset, v uint32) bool {
	if !m.hasLen(offset, 4) {
		return false
	}
	binary.LittleEndian.PutUint32(m.Buffer[offset:], v)
	return true
}

// WriteFloat32Le implements wasm.ModuleContext WriteFloat32Le
func (m *MemoryInstance) WriteFloat32Le(offset uint32, v float32) bool {
	return m.WriteUint32Le(offset, math.Float32bits(v))
}

// WriteUint64Le implements wasm.ModuleContext WriteUint64Le
func (m *MemoryInstance) WriteUint64Le(offset uint32, v uint64) bool {
	if !m.hasLen(offset, 8) {
		return false
	}
	binary.LittleEndian.PutUint64(m.Buffer[offset:], v)
	return true
}

// WriteFloat64Le implements wasm.ModuleContext WriteFloat64Le
func (m *MemoryInstance) WriteFloat64Le(offset uint32, v float64) bool {
	return m.WriteUint64Le(offset, math.Float64bits(v))
}

// Write implements wasm.ModuleContext Write
func (m *MemoryInstance) Write(offset uint32, val []byte) bool {
	if !m.hasLen(offset, uint32(len(val))) {
		return false
	}
	copy(m.Buffer[offset:], val)
	return true
}

// NoopMemory is used when there is no memory or it has a max of zero pages.
var NoopMemory = &noopMemory{}

type noopMemory struct {
}

// Len implements wasm.ModuleContext Len
func (m *noopMemory) Len() uint32 {
	return 0
}

// ReadUint32Le implements wasm.ModuleContext ReadUint32Le
func (m *noopMemory) ReadUint32Le(_ uint32) (uint32, bool) {
	return 0, false
}

// ReadFloat32Le implements wasm.ModuleContext ReadFloat32Le
func (m *noopMemory) ReadFloat32Le(_ uint32) (float32, bool) {
	return 0, false
}

// ReadUint64Le implements wasm.ModuleContext ReadUint64Le
func (m *noopMemory) ReadUint64Le(_ uint32) (uint64, bool) {
	return 0, false
}

// ReadFloat64Le implements wasm.ModuleContext ReadFloat64Le
func (m *noopMemory) ReadFloat64Le(_ uint32) (float64, bool) {
	return 0, false
}

// Read implements wasm.ModuleContext Read
func (m *noopMemory) Read(_, _ uint32) ([]byte, bool) {
	return nil, false
}

// WriteUint32Le implements wasm.ModuleContext WriteUint32Le
func (m *noopMemory) WriteUint32Le(_, _ uint32) bool {
	return false
}

// WriteFloat32Le implements wasm.ModuleContext WriteFloat32Le
func (m *noopMemory) WriteFloat32Le(_ uint32, _ float32) bool {
	return false
}

// WriteUint64Le implements wasm.ModuleContext WriteUint64Le
func (m *noopMemory) WriteUint64Le(_ uint32, _ uint64) bool {
	return false
}

// WriteFloat64Le implements wasm.ModuleContext WriteFloat64Le
func (m *noopMemory) WriteFloat64Le(_ uint32, _ float64) bool {
	return false
}

// Write implements wasm.ModuleContext Write
func (m *noopMemory) Write(_ uint32, _ []byte) bool {
	return false
}