Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat!: add ability to configure size of the Extism var store #61

Merged
merged 5 commits into from
Mar 7, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 42 additions & 11 deletions extism.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"`
Expand All @@ -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"`
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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{
Expand All @@ -435,6 +465,7 @@ func NewPlugin(
LastStatusCode: 0,
Timeout: time.Duration(manifest.Timeout) * time.Millisecond,
MaxHttpResponseBytes: httpMax,
MaxVarBytes: varMax,
log: logStd,
logLevel: logLevel,
}
Expand Down
18 changes: 18 additions & 0 deletions extism_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
38 changes: 23 additions & 15 deletions host.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"io"
"net/http"
"net/url"
"unsafe"

// TODO: is there a better package for this?
"github.com/gobwas/glob"
Expand Down Expand Up @@ -414,34 +415,41 @@ 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)
if err != nil {
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 {
Expand Down
Loading