working commit
This commit is contained in:
+421
@@ -0,0 +1,421 @@
|
||||
// new
|
||||
package extism
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
observe "github.com/dylibso/observe-sdk/go"
|
||||
"github.com/tetratelabs/wazero"
|
||||
"github.com/tetratelabs/wazero/api"
|
||||
"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
|
||||
)
|
||||
|
||||
type CompiledPlugin struct {
|
||||
runtime wazero.Runtime
|
||||
main wazero.CompiledModule
|
||||
extism wazero.CompiledModule
|
||||
env api.Module
|
||||
modules map[string]wazero.CompiledModule
|
||||
|
||||
// when a module (main) is instantiated, it may have a module name that's added
|
||||
// to the data section of the wasm. If this is the case, we won't be able to
|
||||
// instantiate that module more than once. This counter acts as the module name
|
||||
// incrementing each time we instantiate the module.
|
||||
instanceCount atomic.Uint64
|
||||
|
||||
// this is the raw wasm bytes of the provided module, it is required when using a tracing observeAdapter.
|
||||
// If an adapter is not provided, this field will be nil.
|
||||
wasmBytes []byte
|
||||
hasWasi bool
|
||||
manifest Manifest
|
||||
observeAdapter *observe.AdapterBase
|
||||
observeOptions *observe.Options
|
||||
|
||||
maxHttp int64
|
||||
maxVar int64
|
||||
enableHttpResponseHeaders bool
|
||||
}
|
||||
|
||||
type PluginConfig struct {
|
||||
RuntimeConfig wazero.RuntimeConfig
|
||||
EnableWasi bool
|
||||
ObserveAdapter *observe.AdapterBase
|
||||
ObserveOptions *observe.Options
|
||||
EnableHttpResponseHeaders bool
|
||||
|
||||
// ModuleConfig is only used when a plugins are built using the NewPlugin
|
||||
// function. In this function, the plugin is both compiled, and an instance
|
||||
// of the plugin is instantiated, and the ModuleConfig is passed to the
|
||||
// instance.
|
||||
//
|
||||
// When plugins are built using NewCompiledPlugin, the ModuleConfig has no
|
||||
// effect because the instance is not created. Instead, the ModuleConfig is
|
||||
// passed directly in calls to the CompiledPlugin.Instance method.
|
||||
//
|
||||
// NOTE: Module name and start functions are ignored as they are overridden by Extism, also if Manifest contains
|
||||
// non-empty AllowedPaths, then FS is also ignored. If EXTISM_ENABLE_WASI_OUTPUT is set, then stdout and stderr are
|
||||
// set to os.Stdout and os.Stderr respectively (ignoring user defined module config).
|
||||
ModuleConfig wazero.ModuleConfig
|
||||
}
|
||||
|
||||
// NewPlugin creates compiles and instantiates a plugin that is ready
|
||||
// to be used. Plugins are not thread-safe. If you need to use a plugin
|
||||
// across multiple goroutines, use NewCompiledPlugin and create instances
|
||||
// of the plugin using the CompiledPlugin.Instance method.
|
||||
func NewPlugin(
|
||||
ctx context.Context,
|
||||
manifest Manifest,
|
||||
config PluginConfig,
|
||||
functions []HostFunction,
|
||||
) (*Plugin, error) {
|
||||
c, err := NewCompiledPlugin(ctx, manifest, config, functions)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p, err := c.Instance(ctx, PluginInstanceConfig{
|
||||
ModuleConfig: config.ModuleConfig,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p.close = append(p.close, c.Close)
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func calculateMaxHttp(manifest Manifest) int64 {
|
||||
// Default is 50MB
|
||||
maxHttp := int64(1024 * 1024 * 50)
|
||||
if manifest.Memory != nil && manifest.Memory.MaxHttpResponseBytes >= 0 {
|
||||
maxHttp = manifest.Memory.MaxHttpResponseBytes
|
||||
}
|
||||
return maxHttp
|
||||
}
|
||||
|
||||
func calculateMaxVar(manifest Manifest) int64 {
|
||||
// Default is 1MB
|
||||
maxVar := int64(1024 * 1024)
|
||||
if manifest.Memory != nil && manifest.Memory.MaxVarBytes >= 0 {
|
||||
maxVar = manifest.Memory.MaxVarBytes
|
||||
}
|
||||
return maxVar
|
||||
}
|
||||
|
||||
// NewCompiledPlugin creates a compiled plugin that is ready to be instantiated.
|
||||
// You can instantiate the plugin multiple times using the CompiledPlugin.Instance
|
||||
// method and run those instances concurrently.
|
||||
func NewCompiledPlugin(
|
||||
ctx context.Context,
|
||||
manifest Manifest,
|
||||
config PluginConfig,
|
||||
funcs []HostFunction,
|
||||
) (*CompiledPlugin, error) {
|
||||
count := len(manifest.Wasm)
|
||||
if count == 0 {
|
||||
return nil, fmt.Errorf("manifest can't be empty")
|
||||
}
|
||||
|
||||
runtimeConfig := config.RuntimeConfig
|
||||
if runtimeConfig == nil {
|
||||
runtimeConfig = wazero.NewRuntimeConfig()
|
||||
}
|
||||
|
||||
// Make sure function calls are cancelled if the context is cancelled
|
||||
if manifest.Timeout > 0 {
|
||||
runtimeConfig = runtimeConfig.WithCloseOnContextDone(true)
|
||||
}
|
||||
|
||||
if manifest.Memory != nil {
|
||||
if manifest.Memory.MaxPages > 0 {
|
||||
runtimeConfig = runtimeConfig.WithMemoryLimitPages(manifest.Memory.MaxPages)
|
||||
}
|
||||
}
|
||||
|
||||
p := CompiledPlugin{
|
||||
manifest: manifest,
|
||||
runtime: wazero.NewRuntimeWithConfig(ctx, runtimeConfig),
|
||||
observeAdapter: config.ObserveAdapter,
|
||||
observeOptions: config.ObserveOptions,
|
||||
enableHttpResponseHeaders: config.EnableHttpResponseHeaders,
|
||||
modules: make(map[string]wazero.CompiledModule),
|
||||
maxHttp: calculateMaxHttp(manifest),
|
||||
maxVar: calculateMaxVar(manifest),
|
||||
}
|
||||
|
||||
if config.EnableWasi {
|
||||
wasi_snapshot_preview1.MustInstantiate(ctx, p.runtime)
|
||||
p.hasWasi = true
|
||||
}
|
||||
|
||||
// Build host modules
|
||||
hostModules := make(map[string][]HostFunction)
|
||||
for _, f := range funcs {
|
||||
hostModules[f.Namespace] = append(hostModules[f.Namespace], f)
|
||||
}
|
||||
for name, funcs := range hostModules {
|
||||
_, err := buildHostModule(ctx, p.runtime, name, funcs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("building host module: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Compile the extism module
|
||||
var err error
|
||||
p.extism, err = p.runtime.CompileModule(ctx, extismRuntimeWasm)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("instantiating extism module: %w", err)
|
||||
}
|
||||
|
||||
// Build and instantiate extism:host/env module
|
||||
p.env, err = instantiateEnvModule(ctx, p.runtime)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Try to find the main module:
|
||||
// - There is always one main module
|
||||
// - If a Wasm value has the Name field set to "main" then use that module
|
||||
// - If there is only one module in the manifest then that is the main module by default
|
||||
// - Otherwise the last module listed is the main module
|
||||
|
||||
foundMain := false
|
||||
for i, wasm := range manifest.Wasm {
|
||||
data, err := wasm.ToWasmData(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if (data.Name == "" || i == len(manifest.Wasm)-1) && !foundMain {
|
||||
data.Name = "main"
|
||||
}
|
||||
|
||||
_, okm := p.modules[data.Name]
|
||||
|
||||
if data.Name == "extism:host/env" || okm {
|
||||
return nil, fmt.Errorf("module name collision: '%s'", data.Name)
|
||||
}
|
||||
|
||||
if data.Hash != "" {
|
||||
calculatedHash := calculateHash(data.Data)
|
||||
if data.Hash != calculatedHash {
|
||||
return nil, fmt.Errorf("hash mismatch for module '%s'", data.Name)
|
||||
}
|
||||
}
|
||||
|
||||
if p.observeAdapter != nil {
|
||||
p.wasmBytes = data.Data
|
||||
}
|
||||
|
||||
compiledModule, err := p.runtime.CompileModule(ctx, data.Data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if data.Name == "main" {
|
||||
if foundMain {
|
||||
return nil, errors.New("can't have more than one main module")
|
||||
}
|
||||
p.main = compiledModule
|
||||
foundMain = true
|
||||
} else {
|
||||
// Store compiled module for instantiation
|
||||
p.modules[data.Name] = compiledModule
|
||||
// Create wrapper with original name that will forward calls to the actual module instance. See createModuleWrapper for more details.
|
||||
_, err = createModuleWrapper(ctx, p.runtime, data.Name, compiledModule)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create wrapper for %s: %w", data.Name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if p.main == nil {
|
||||
return nil, errors.New("no main module found")
|
||||
}
|
||||
|
||||
// We no longer need the wasm in the manifest so nil it
|
||||
// to make the slice eligible for garbage collection.
|
||||
p.manifest.Wasm = nil
|
||||
|
||||
return &p, nil
|
||||
}
|
||||
|
||||
// createModuleWrapper creates a host module that acts as a proxy for module instances.
|
||||
// In Wazero, modules with the same name cannot be instantiated multiple times in the same runtime.
|
||||
// However, we need each Plugin instance to have its own copy of each module for isolation. To solve this, we:
|
||||
// 1. Create a host module wrapper that keeps the original module name (needed for imports to work)
|
||||
// 2. Instantiate actual module copies with unique names for each Plugin
|
||||
// 3. The wrapper forwards function calls to the correct module instance for each Plugin
|
||||
func createModuleWrapper(ctx context.Context, rt wazero.Runtime, name string, compiled wazero.CompiledModule) (api.Module, error) {
|
||||
builder := rt.NewHostModuleBuilder(name)
|
||||
|
||||
// Create proxy functions for each exported function from the original module.
|
||||
// These proxies will forward calls to the appropriate module instance.
|
||||
for _, export := range compiled.ExportedFunctions() {
|
||||
exportName := export.Name()
|
||||
|
||||
// Skip wrapping the _start function since it's automatically called by wazero during instantiation.
|
||||
// The wrapper functions require a Plugin instance in the context to work, but during wrapper
|
||||
// instantiation there is no Plugin instance yet.
|
||||
if exportName == "_start" {
|
||||
continue
|
||||
}
|
||||
|
||||
// Create a proxy function that:
|
||||
// 1. Gets the calling Plugin instance from context
|
||||
// 2. Looks up that Plugin's copy of this module
|
||||
// 3. Forwards the call to the actual function
|
||||
wrapper := func(callCtx context.Context, mod api.Module, stack []uint64) {
|
||||
// Get the Plugin instance that's making this call
|
||||
plugin, ok := callCtx.Value(PluginCtxKey("plugin")).(*Plugin)
|
||||
if !ok {
|
||||
panic("Invalid context, `plugin` key not found")
|
||||
}
|
||||
|
||||
// Get this Plugin's instance of the module
|
||||
actualModule, ok := plugin.modules[name]
|
||||
if !ok {
|
||||
panic(fmt.Sprintf("module %s not found in plugin", name))
|
||||
}
|
||||
|
||||
// Forward the call to the actual module instance
|
||||
fn := actualModule.ExportedFunction(exportName)
|
||||
if fn == nil {
|
||||
panic(fmt.Sprintf("function %s not found in module %s", exportName, name))
|
||||
}
|
||||
|
||||
err := fn.CallWithStack(callCtx, stack)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Export the proxy function with the same name and signature as the original
|
||||
builder.NewFunctionBuilder().
|
||||
WithGoModuleFunction(api.GoModuleFunc(wrapper), export.ParamTypes(), export.ResultTypes()).
|
||||
Export(exportName)
|
||||
}
|
||||
|
||||
return builder.Instantiate(ctx)
|
||||
}
|
||||
|
||||
func (p *CompiledPlugin) Close(ctx context.Context) error {
|
||||
return p.runtime.Close(ctx)
|
||||
}
|
||||
|
||||
func (p *CompiledPlugin) Instance(ctx context.Context, config PluginInstanceConfig) (*Plugin, error) {
|
||||
instanceNum := p.instanceCount.Add(1)
|
||||
|
||||
var closers []func(ctx context.Context) error
|
||||
|
||||
moduleConfig := config.ModuleConfig
|
||||
if moduleConfig == nil {
|
||||
moduleConfig = wazero.NewModuleConfig()
|
||||
}
|
||||
|
||||
// NOTE: we don't want wazero to call the start function, we will initialize
|
||||
// the guest runtime manually.
|
||||
// See: https://github.com/extism/go-sdk/pull/1#issuecomment-1650527495
|
||||
moduleConfig = moduleConfig.WithStartFunctions()
|
||||
|
||||
if len(p.manifest.AllowedPaths) > 0 {
|
||||
// NOTE: this is only necessary for guest modules because
|
||||
// host modules have the same access privileges as the host itself
|
||||
fs := wazero.NewFSConfig()
|
||||
for host, guest := range p.manifest.AllowedPaths {
|
||||
if strings.HasPrefix(host, "ro:") {
|
||||
trimmed := strings.TrimPrefix(host, "ro:")
|
||||
fs = fs.WithReadOnlyDirMount(trimmed, guest)
|
||||
} else {
|
||||
fs = fs.WithDirMount(host, guest)
|
||||
}
|
||||
}
|
||||
moduleConfig = moduleConfig.WithFSConfig(fs)
|
||||
}
|
||||
|
||||
_, wasiOutput := os.LookupEnv("EXTISM_ENABLE_WASI_OUTPUT")
|
||||
if p.hasWasi && wasiOutput {
|
||||
moduleConfig = moduleConfig.WithStderr(os.Stderr).WithStdout(os.Stdout)
|
||||
}
|
||||
|
||||
var trace *observe.TraceCtx
|
||||
var err error
|
||||
if p.observeAdapter != nil {
|
||||
trace, err = p.observeAdapter.NewTraceCtx(ctx, p.runtime, p.wasmBytes, p.observeOptions)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to initialize Observe Adapter: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Compile and instantiate the extism runtime. This runtime is stateful and needs to be
|
||||
// instantiated on a per-instance basis. We don't provide a name because the module needs
|
||||
// to be anonymous -- you cannot instantiate multiple modules with the same name into the
|
||||
// same runtime. It is okay that this is anonymous, because this module is only called
|
||||
// from Go host functions and not from the Wasm module itself.
|
||||
extism, err := p.runtime.InstantiateModule(ctx, p.extism, wazero.NewModuleConfig())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("instantiating extism module: %w", err)
|
||||
}
|
||||
|
||||
closers = append(closers, extism.Close)
|
||||
|
||||
// Instantiate all non-main modules first
|
||||
instancedModules := make(map[string]api.Module)
|
||||
for name, compiledModule := range p.modules {
|
||||
uniqueName := fmt.Sprintf("%s_%d", name, instanceNum)
|
||||
instance, err := p.runtime.InstantiateModule(ctx, compiledModule, moduleConfig.WithName(uniqueName))
|
||||
if err != nil {
|
||||
for _, closer := range closers {
|
||||
closer(ctx)
|
||||
}
|
||||
return nil, fmt.Errorf("instantiating module %s: %w", name, err)
|
||||
}
|
||||
instancedModules[name] = instance
|
||||
closers = append(closers, instance.Close)
|
||||
}
|
||||
|
||||
mainModuleName := fmt.Sprintf("main_%d", instanceNum)
|
||||
main, err := p.runtime.InstantiateModule(ctx, p.main, moduleConfig.WithName(mainModuleName))
|
||||
if err != nil {
|
||||
for _, closer := range closers {
|
||||
closer(ctx)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("instantiating module: %w", err)
|
||||
}
|
||||
|
||||
closers = append(closers, main.Close)
|
||||
|
||||
var headers map[string]string = nil
|
||||
if p.enableHttpResponseHeaders {
|
||||
headers = map[string]string{}
|
||||
}
|
||||
|
||||
instance := &Plugin{
|
||||
close: closers,
|
||||
extism: extism,
|
||||
hasWasi: p.hasWasi,
|
||||
mainModule: main,
|
||||
modules: instancedModules,
|
||||
Timeout: time.Duration(p.manifest.Timeout) * time.Millisecond,
|
||||
Config: p.manifest.Config,
|
||||
Var: make(map[string][]byte),
|
||||
AllowedHosts: p.manifest.AllowedHosts,
|
||||
AllowedPaths: p.manifest.AllowedPaths,
|
||||
LastStatusCode: 0,
|
||||
LastResponseHeaders: headers,
|
||||
MaxHttpResponseBytes: p.maxHttp,
|
||||
MaxVarBytes: p.maxVar,
|
||||
guestRuntime: guestRuntime{},
|
||||
Adapter: p.observeAdapter,
|
||||
log: logStd,
|
||||
traceCtx: trace,
|
||||
}
|
||||
instance.guestRuntime = detectGuestRuntime(instance)
|
||||
|
||||
return instance, nil
|
||||
}
|
||||
Reference in New Issue
Block a user