From a1a281594a4e43ebddce1898d90128ea86191ada Mon Sep 17 00:00:00 2001 From: zach Date: Thu, 7 Mar 2024 09:55:26 -0800 Subject: [PATCH] feat!: add ability to configure size of the Extism var store (#61) See https://github.com/extism/extism/pull/692 - Breaking: changes the `Memory` field of `Manifest` to a pointer in order to better detect when the memory settings weren't provided --- extism.go | 53 +++++++++++++++++++++++++++++++++++++++----------- extism_test.go | 18 +++++++++++++++++ host.go | 38 ++++++++++++++++++++++-------------- 3 files changed, 83 insertions(+), 26 deletions(-) diff --git a/extism.go b/extism.go index 9eb902b..03c57e1 100644 --- a/extism.go +++ b/extism.go @@ -90,6 +90,7 @@ type Plugin struct { AllowedPaths map[string]string LastStatusCode int MaxHttpResponseBytes int64 + MaxVarBytes int64 log func(LogLevel, string) logLevel LogLevel guestRuntime guestRuntime @@ -216,14 +217,17 @@ func (u WasmUrl) ToWasmData(ctx context.Context) (WasmData, error) { }, nil } +type ManifestMemory struct { + MaxPages uint32 `json:"max_pages,omitempty"` + MaxHttpResponseBytes int64 `json:"max_http_response_bytes,omitempty"` + MaxVarBytes int64 `json:"max_var_bytes,omitempty"` +} + // Manifest represents the plugin's manifest, including Wasm modules and configuration. // See https://extism.org/docs/concepts/manifest for schema. type Manifest struct { - Wasm []Wasm `json:"wasm"` - Memory struct { - MaxPages uint32 `json:"max_pages,omitempty"` - MaxHttpResponseBytes uint64 `json:"max_http_response_bytes,omitempty"` - } `json:"memory,omitempty"` + Wasm []Wasm `json:"wasm"` + Memory *ManifestMemory `json:"memory,omitempty"` Config map[string]string `json:"config,omitempty"` AllowedHosts []string `json:"allowed_hosts,omitempty"` AllowedPaths map[string]string `json:"allowed_paths,omitempty"` @@ -232,9 +236,10 @@ type Manifest struct { type concreteManifest struct { Wasm []concreteWasm `json:"wasm"` - Memory struct { + Memory *struct { MaxPages uint32 `json:"max_pages,omitempty"` - MaxHttpResponseBytes uint64 `json:"max_http_response_bytes,omitempty"` + MaxHttpResponseBytes *int64 `json:"max_http_response_bytes,omitempty"` + MaxVarBytes *int64 `json:"max_var_bytes,omitempty"` } `json:"memory,omitempty"` Config map[string]string `json:"config,omitempty"` AllowedHosts []string `json:"allowed_hosts,omitempty"` @@ -249,7 +254,25 @@ func (m *Manifest) UnmarshalJSON(data []byte) error { return err } - m.Memory = tmp.Memory + m.Memory = &ManifestMemory{} + if tmp.Memory != nil { + m.Memory.MaxPages = tmp.Memory.MaxPages + if tmp.Memory.MaxHttpResponseBytes != nil { + m.Memory.MaxHttpResponseBytes = *tmp.Memory.MaxHttpResponseBytes + } else { + m.Memory.MaxHttpResponseBytes = -1 + } + + if tmp.Memory.MaxVarBytes != nil { + m.Memory.MaxVarBytes = *tmp.Memory.MaxVarBytes + } else { + m.Memory.MaxVarBytes = -1 + } + } else { + m.Memory.MaxPages = 0 + m.Memory.MaxHttpResponseBytes = -1 + m.Memory.MaxVarBytes = -1 + } m.Config = tmp.Config m.AllowedHosts = tmp.AllowedHosts m.AllowedPaths = tmp.AllowedPaths @@ -301,8 +324,10 @@ func NewPlugin( rconfig = rconfig.WithCloseOnContextDone(true) } - if manifest.Memory.MaxPages > 0 { - rconfig = rconfig.WithMemoryLimitPages(manifest.Memory.MaxPages) + if manifest.Memory != nil { + if manifest.Memory.MaxPages > 0 { + rconfig = rconfig.WithMemoryLimitPages(manifest.Memory.MaxPages) + } } rt := wazero.NewRuntimeWithConfig(ctx, rconfig) @@ -419,9 +444,14 @@ func NewPlugin( i := 0 httpMax := int64(1024 * 1024 * 50) - if manifest.Memory.MaxHttpResponseBytes != 0 { + if manifest.Memory != nil && manifest.Memory.MaxHttpResponseBytes >= 0 { httpMax = int64(manifest.Memory.MaxHttpResponseBytes) } + + varMax := int64(1024 * 1024) + if manifest.Memory != nil && manifest.Memory.MaxVarBytes >= 0 { + varMax = int64(manifest.Memory.MaxVarBytes) + } for _, m := range modules { if m.Name() == "main" { p := &Plugin{ @@ -435,6 +465,7 @@ func NewPlugin( LastStatusCode: 0, Timeout: time.Duration(manifest.Timeout) * time.Millisecond, MaxHttpResponseBytes: httpMax, + MaxVarBytes: varMax, log: logStd, logLevel: logLevel, } diff --git a/extism_test.go b/extism_test.go index bb4ee75..d2349fe 100644 --- a/extism_test.go +++ b/extism_test.go @@ -573,6 +573,24 @@ func TestVar(t *testing.T) { } +func TestNoVars(t *testing.T) { + manifest := manifest("var.wasm") + manifest.Memory = &ManifestMemory{MaxVarBytes: 0} + + if plugin, ok := plugin(t, manifest); ok { + defer plugin.Close() + + plugin.Var["a"] = uintToLEBytes(10) + + _, _, err := plugin.Call("run_test", []byte{}) + + if err == nil { + t.Fail() + } + } + +} + func TestFS(t *testing.T) { manifest := manifest("fs.wasm") manifest.AllowedPaths = map[string]string{ diff --git a/host.go b/host.go index bcc4148..e28e444 100644 --- a/host.go +++ b/host.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "net/url" + "unsafe" // TODO: is there a better package for this? "github.com/gobwas/glob" @@ -414,6 +415,10 @@ func varSet(ctx context.Context, m api.Module, nameOffset uint64, valueOffset ui panic("Invalid context, `plugin` key not found") } + if plugin.MaxVarBytes == 0 { + panic("Vars are disabled by this host") + } + cp := plugin.currentPlugin() name, err := cp.ReadString(nameOffset) @@ -421,27 +426,30 @@ func varSet(ctx context.Context, m api.Module, nameOffset uint64, valueOffset ui panic(fmt.Errorf("Failed to read var name from memory: %v", err)) } - size := 0 - for _, v := range plugin.Var { - size += len(v) + // Remove if the value offset is 0 + if valueOffset == 0 { + delete(plugin.Var, name) + return } - // If the store is larger than 100MB then stop adding things - if size > 1024*1024*100 && valueOffset != 0 { - panic("Variable store is full") + value, err := cp.ReadBytes(valueOffset) + if err != nil { + panic(fmt.Errorf("Failed to read var value from memory: %v", err)) } - // Remove if the value offset is 0 - if valueOffset == 0 { - delete(plugin.Var, name) - } else { - value, err := cp.ReadBytes(valueOffset) - if err != nil { - panic(fmt.Errorf("Failed to read var value from memory: %v", err)) - } + // Calculate size including current key/value + size := int(unsafe.Sizeof([]byte{})+unsafe.Sizeof("")) + len(name) + len(value) + for k, v := range plugin.Var { + size += len(k) + size += len(v) + size += int(unsafe.Sizeof([]byte{}) + unsafe.Sizeof("")) + } - plugin.Var[name] = value + if size >= int(plugin.MaxVarBytes) && valueOffset != 0 { + panic("Variable store is full") } + + plugin.Var[name] = value } func httpRequest(ctx context.Context, m api.Module, requestOffset uint64, bodyOffset uint64) uint64 {