Enhance wasm plugins

Co-authored-by: Michael <[michael.matur@gmail.com](mailto:michael.matur@gmail.com)>

pkg/plugins/middlewarewasm.go

@@ -8,20 +8,38 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"strings"
 
 	"github.com/http-wasm/http-wasm-host-go/handler"
 	wasm "github.com/http-wasm/http-wasm-host-go/handler/nethttp"
+	"github.com/juliens/wasm-goexport/host"
 	"github.com/tetratelabs/wazero"
 	"github.com/traefik/traefik/v3/pkg/logs"
 	"github.com/traefik/traefik/v3/pkg/middlewares"
 )
 
 type wasmMiddlewareBuilder struct {
-	path string
+	path     string
+	cache    wazero.CompilationCache
+	settings Settings
 }
 
-func newWasmMiddlewareBuilder(goPath string, moduleName, wasmPath string) *wasmMiddlewareBuilder {
-	return &wasmMiddlewareBuilder{path: filepath.Join(goPath, "src", moduleName, wasmPath)}
+func newWasmMiddlewareBuilder(goPath, moduleName, wasmPath string, settings Settings) (*wasmMiddlewareBuilder, error) {
+	ctx := context.Background()
+	path := filepath.Join(goPath, "src", moduleName, wasmPath)
+	cache := wazero.NewCompilationCache()
+
+	code, err := os.ReadFile(path)
+	if err != nil {
+		return nil, fmt.Errorf("loading Wasm binary: %w", err)
+	}
+
+	rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().WithCompilationCache(cache))
+	if _, err = rt.CompileModule(ctx, code); err != nil {
+		return nil, fmt.Errorf("compiling guest module: %w", err)
+	}
+
+	return &wasmMiddlewareBuilder{path: path, cache: cache, settings: settings}, nil
 }
 
 func (b wasmMiddlewareBuilder) newMiddleware(config map[string]interface{}, middlewareName string) (pluginMiddleware, error) {
@@ -33,15 +51,64 @@ func (b wasmMiddlewareBuilder) newMiddleware(config map[string]interface{}, midd
 }
 
 func (b wasmMiddlewareBuilder) newHandler(ctx context.Context, next http.Handler, cfg reflect.Value, middlewareName string) (http.Handler, error) {
+	h, applyCtx, err := b.buildMiddleware(ctx, next, cfg, middlewareName)
+	if err != nil {
+		return nil, fmt.Errorf("building Wasm middleware: %w", err)
+	}
+
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		h.ServeHTTP(rw, req.WithContext(applyCtx(req.Context())))
+	}), nil
+}
+
+func (b *wasmMiddlewareBuilder) buildMiddleware(ctx context.Context, next http.Handler, cfg reflect.Value, middlewareName string) (http.Handler, func(ctx context.Context) context.Context, error) {
 	code, err := os.ReadFile(b.path)
 	if err != nil {
-		return nil, fmt.Errorf("loading Wasm binary: %w", err)
+		return nil, nil, fmt.Errorf("loading binary: %w", err)
+	}
+
+	rt := host.NewRuntime(wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().WithCompilationCache(b.cache)))
+
+	guestModule, err := rt.CompileModule(ctx, code)
+	if err != nil {
+		return nil, nil, fmt.Errorf("compiling guest module: %w", err)
+	}
+
+	applyCtx, err := InstantiateHost(ctx, rt, guestModule, b.settings)
+	if err != nil {
+		return nil, nil, fmt.Errorf("instantiating host module: %w", err)
 	}
 
 	logger := middlewares.GetLogger(ctx, middlewareName, "wasm")
 
+	config := wazero.NewModuleConfig().WithSysWalltime()
+	for _, env := range b.settings.Envs {
+		config.WithEnv(env, os.Getenv(env))
+	}
+
+	if len(b.settings.Mounts) > 0 {
+		fsConfig := wazero.NewFSConfig()
+		for _, mount := range b.settings.Mounts {
+			withDir := fsConfig.WithDirMount
+			prefix, readOnly := strings.CutSuffix(mount, ":ro")
+			if readOnly {
+				withDir = fsConfig.WithReadOnlyDirMount
+			}
+			parts := strings.Split(prefix, ":")
+			switch {
+			case len(parts) == 1:
+				withDir(parts[0], parts[0])
+			case len(parts) == 2:
+				withDir(parts[0], parts[1])
+			default:
+				return nil, nil, fmt.Errorf("invalid directory %q", mount)
+			}
+		}
+		config.WithFSConfig(fsConfig)
+	}
+
 	opts := []handler.Option{
-		handler.ModuleConfig(wazero.NewModuleConfig().WithSysWalltime()),
+		handler.ModuleConfig(config),
 		handler.Logger(logs.NewWasmLogger(logger)),
 	}
 
@@ -49,23 +116,27 @@ func (b wasmMiddlewareBuilder) newHandler(ctx context.Context, next http.Handler
 	if i != nil {
 		config, ok := i.(map[string]interface{})
 		if !ok {
-			return nil, fmt.Errorf("could not type assert config: %T", i)
+			return nil, nil, fmt.Errorf("could not type assert config: %T", i)
 		}
 
 		data, err := json.Marshal(config)
 		if err != nil {
-			return nil, fmt.Errorf("marshaling config: %w", err)
+			return nil, nil, fmt.Errorf("marshaling config: %w", err)
 		}
 
 		opts = append(opts, handler.GuestConfig(data))
 	}
 
-	mw, err := wasm.NewMiddleware(context.Background(), code, opts...)
+	opts = append(opts, handler.Runtime(func(ctx context.Context) (wazero.Runtime, error) {
+		return rt, nil
+	}))
+
+	mw, err := wasm.NewMiddleware(applyCtx(ctx), code, opts...)
 	if err != nil {
-		return nil, err
+		return nil, nil, fmt.Errorf("creating middleware: %w", err)
 	}
 
-	return mw.NewHandler(ctx, next), nil
+	return mw.NewHandler(ctx, next), applyCtx, nil
 }
 
 // WasmMiddleware is an HTTP handler plugin wrapper.