traefik/pkg/plugins/middlewarewasm.go

152 lines
4.4 KiB
Go
Raw Normal View History

package plugins
import (
"context"
"encoding/json"
"fmt"
"net/http"
"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/tetratelabs/wazero"
"github.com/traefik/traefik/v3/pkg/logs"
"github.com/traefik/traefik/v3/pkg/middlewares"
)
type wasmMiddlewareBuilder struct {
path string
cache wazero.CompilationCache
settings Settings
}
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) {
return &WasmMiddleware{
middlewareName: middlewareName,
config: reflect.ValueOf(config),
builder: b,
}, nil
}
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, nil, fmt.Errorf("loading binary: %w", err)
}
rt := 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().WithStartFunctions("_start", "_initialize")
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(config),
handler.Logger(logs.NewWasmLogger(logger)),
}
i := cfg.Interface()
if i != nil {
config, ok := i.(map[string]interface{})
if !ok {
return nil, nil, fmt.Errorf("could not type assert config: %T", i)
}
data, err := json.Marshal(config)
if err != nil {
return nil, nil, fmt.Errorf("marshaling config: %w", err)
}
opts = append(opts, handler.GuestConfig(data))
}
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, nil, fmt.Errorf("creating middleware: %w", err)
}
return mw.NewHandler(ctx, next), applyCtx, nil
}
// WasmMiddleware is an HTTP handler plugin wrapper.
type WasmMiddleware struct {
middlewareName string
config reflect.Value
builder wasmMiddlewareBuilder
}
// NewHandler creates a new HTTP handler.
func (m WasmMiddleware) NewHandler(ctx context.Context, next http.Handler) (http.Handler, error) {
return m.builder.newHandler(ctx, next, m.config, m.middlewareName)
}