runtime.go

Documentation: cuelang.org/go/cue/interpreter/wasm

     1  // Copyright 2023 CUE Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package wasm
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"os"
    21  	"sync"
    22  
    23  	"github.com/tetratelabs/wazero"
    24  	"github.com/tetratelabs/wazero/api"
    25  	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
    26  )
    27  
    28  // A runtime is a Wasm runtime that can compile, load, and execute
    29  // Wasm code.
    30  type runtime struct {
    31  	// ctx exists so that we have something to pass to Wazero
    32  	// functions, but it's unused otherwise.
    33  	ctx context.Context
    34  
    35  	wazero.Runtime
    36  }
    37  
    38  func newRuntime() runtime {
    39  	ctx := context.Background()
    40  	r := wazero.NewRuntime(ctx)
    41  	wasi_snapshot_preview1.MustInstantiate(ctx, r)
    42  
    43  	return runtime{
    44  		ctx:     ctx,
    45  		Runtime: r,
    46  	}
    47  }
    48  
    49  // compile takes the name of a Wasm module, and returns its compiled
    50  // form, or an error.
    51  func (r *runtime) compile(name string) (*module, error) {
    52  	buf, err := os.ReadFile(name)
    53  	if err != nil {
    54  		return nil, fmt.Errorf("can't compile Wasm module: %w", err)
    55  	}
    56  
    57  	mod, err := r.Runtime.CompileModule(r.ctx, buf)
    58  	if err != nil {
    59  		return nil, fmt.Errorf("can't compile Wasm module: %w", err)
    60  	}
    61  	return &module{
    62  		runtime:        r,
    63  		name:           name,
    64  		CompiledModule: mod,
    65  	}, nil
    66  }
    67  
    68  // compileAndLoad is a convenience method that compiles a module then
    69  // loads it into memory returning the loaded instance, or an error.
    70  func (r *runtime) compileAndLoad(name string) (*instance, error) {
    71  	m, err := r.compile(name)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  	i, err := m.load()
    76  	if err != nil {
    77  		return nil, err
    78  	}
    79  	return i, nil
    80  }
    81  
    82  // A module is a compiled Wasm module.
    83  type module struct {
    84  	*runtime
    85  	name string
    86  	wazero.CompiledModule
    87  }
    88  
    89  // load loads the compiled module into memory, returning a new instance
    90  // that can be called into, or an error. Different instances of the
    91  // same module do not share memory.
    92  func (m *module) load() (*instance, error) {
    93  	cfg := wazero.NewModuleConfig().WithName(m.name)
    94  	wInst, err := m.Runtime.InstantiateModule(m.ctx, m.CompiledModule, cfg)
    95  	if err != nil {
    96  		return nil, fmt.Errorf("can't instantiate Wasm module: %w", err)
    97  	}
    98  
    99  	inst := instance{
   100  		module:   m,
   101  		instance: wInst,
   102  		alloc:    wInst.ExportedFunction("allocate"),
   103  		free:     wInst.ExportedFunction("deallocate"),
   104  	}
   105  	return &inst, nil
   106  }
   107  
   108  // An instance is a Wasm module loaded into memory.
   109  type instance struct {
   110  	// mu serializes access the whole struct.
   111  	mu sync.Mutex
   112  
   113  	*module
   114  	instance api.Module
   115  
   116  	// alloc is a guest function that allocates guest memory on
   117  	// behalf of the host.
   118  	alloc api.Function
   119  
   120  	// free is a guest function that frees guest memory on
   121  	// behalf of the host.
   122  	free api.Function
   123  }
   124  
   125  // load attempts to load the named function from the instance, returning
   126  // it if found, or an error.
   127  func (i *instance) load(funcName string) (api.Function, error) {
   128  	i.mu.Lock()
   129  	defer i.mu.Unlock()
   130  
   131  	f := i.instance.ExportedFunction(funcName)
   132  	if f == nil {
   133  		return nil, fmt.Errorf("can't find function %q in Wasm module %v", funcName, i.module.Name())
   134  	}
   135  	return f, nil
   136  }
   137  
   138  // Alloc returns a reference to newly allocated guest memory that spans
   139  // the provided size.
   140  func (i *instance) Alloc(size uint32) (*memory, error) {
   141  	i.mu.Lock()
   142  	defer i.mu.Unlock()
   143  
   144  	res, err := i.alloc.Call(i.ctx, uint64(size))
   145  	if err != nil {
   146  		return nil, fmt.Errorf("can't allocate memory: requested %d bytes", size)
   147  	}
   148  	return &memory{
   149  		i:   i,
   150  		ptr: uint32(res[0]),
   151  		len: size,
   152  	}, nil
   153  }
   154  
   155  // Free frees previously allocated guest memory.
   156  func (i *instance) Free(m *memory) {
   157  	i.mu.Lock()
   158  	defer i.mu.Unlock()
   159  
   160  	i.free.Call(i.ctx, uint64(m.ptr), uint64(m.len))
   161  }
   162  
   163  // Free frees several previously allocated guest memories.
   164  func (i *instance) FreeAll(ms []*memory) {
   165  	i.mu.Lock()
   166  	defer i.mu.Unlock()
   167  
   168  	for _, m := range ms {
   169  		i.free.Call(i.ctx, uint64(m.ptr), uint64(m.len))
   170  	}
   171  }
   172  
   173  // memory is a read and write reference to guest memory that the host
   174  // requested.
   175  type memory struct {
   176  	i   *instance
   177  	ptr uint32
   178  	len uint32
   179  }
   180  
   181  // Bytes return a copy of the contents of the guest memory to the host.
   182  func (m *memory) Bytes() []byte {
   183  	m.i.mu.Lock()
   184  	defer m.i.mu.Unlock()
   185  
   186  	bytes, ok := m.i.instance.Memory().Read(m.ptr, m.len)
   187  	if !ok {
   188  		panic(fmt.Sprintf("can't read %d bytes from Wasm address %#x", m.len, m.ptr))
   189  	}
   190  	return append([]byte{}, bytes...)
   191  }
   192  
   193  // WriteAt writes p at the given relative offset within m.
   194  // It panics if buf doesn't fit into m, or if off is out of bounds.
   195  func (m *memory) WriteAt(p []byte, off int64) (int, error) {
   196  	if (off < 0) || (off >= 1<<32-1) {
   197  		panic(fmt.Sprintf("can't write %d bytes to Wasm address %#x", len(p), m.ptr))
   198  	}
   199  
   200  	m.i.mu.Lock()
   201  	defer m.i.mu.Unlock()
   202  
   203  	ok := m.i.instance.Memory().Write(m.ptr+uint32(off), p)
   204  	if !ok {
   205  		panic(fmt.Sprintf("can't write %d bytes to Wasm address %#x", len(p), m.ptr))
   206  	}
   207  	return len(p), nil
   208  }
   209  
   210  // Args returns a memory in the form of pair of arguments directly
   211  // passable to Wasm.
   212  func (m *memory) Args() []uint64 {
   213  	return []uint64{uint64(m.ptr), uint64(m.len)}
   214  }
   215
View as plain text