feat: add ability to configure size of the Extism var store (#692)

- Adds `memory.max_var_bytes` to the manifest to limit the number of
bytes allowed to be stored in Extism vars - if `max_var_bytes` is set to
0 then vars are disabled.
- Adds some builder functions to `MemoryOptions` struct
- Sets the default var store size to 1mb
- Includes a test to make sure `var_set` returns an error when the limit
is reached
This commit is contained in:
zach
2024-03-07 09:55:02 -08:00
committed by GitHub
parent 8a29e5b1d4
commit d47af24552
4 changed files with 90 additions and 10 deletions

View File

@@ -38,7 +38,8 @@
"description": "Memory options",
"default": {
"max_http_response_bytes": null,
"max_pages": null
"max_pages": null,
"max_var_bytes": null
},
"allOf": [
{
@@ -89,6 +90,16 @@
],
"format": "uint32",
"minimum": 0.0
},
"max_var_bytes": {
"description": "The maximum number of bytes allowed to be used by plugin vars. Setting this to 0 will disable Extism vars. The default value is 1mb.",
"default": 1048576,
"type": [
"integer",
"null"
],
"format": "uint64",
"minimum": 0.0
}
},
"additionalProperties": false

View File

@@ -16,6 +16,40 @@ pub struct MemoryOptions {
/// The maximum number of bytes allowed in an HTTP response
#[serde(default)]
pub max_http_response_bytes: Option<u64>,
/// The maximum number of bytes allowed to be used by plugin vars. Setting this to 0
/// will disable Extism vars. The default value is 1mb.
#[serde(default = "default_var_bytes")]
pub max_var_bytes: Option<u64>,
}
impl MemoryOptions {
/// Create an empty `MemoryOptions` value
pub fn new() -> Self {
Default::default()
}
/// Set max pages
pub fn with_max_pages(mut self, pages: u32) -> Self {
self.max_pages = Some(pages);
self
}
/// Set max HTTP response size
pub fn with_max_http_response_bytes(mut self, bytes: u64) -> Self {
self.max_http_response_bytes = Some(bytes);
self
}
/// Set max size of Extism vars
pub fn with_max_var_bytes(mut self, bytes: u64) -> Self {
self.max_var_bytes = Some(bytes);
self
}
}
fn default_var_bytes() -> Option<u64> {
Some(1024 * 1024)
}
/// Generic HTTP request structure

View File

@@ -97,19 +97,13 @@ pub(crate) fn var_set(
) -> Result<(), Error> {
let data: &mut CurrentPlugin = caller.data_mut();
let mut size = 0;
for v in data.vars.values() {
size += v.len();
if data.manifest.memory.max_var_bytes.is_some_and(|x| x == 0) {
anyhow::bail!("Vars are disabled by this host")
}
let voffset = args!(input, 1, i64) as u64;
// If the store is larger than 100MB then stop adding things
if size > 1024 * 1024 * 100 && voffset != 0 {
return Err(Error::msg("Variable store is full"));
}
let key_offs = args!(input, 0, i64) as u64;
let key = {
let handle = match data.memory_handle(key_offs) {
Some(h) => h,
@@ -132,6 +126,22 @@ pub(crate) fn var_set(
None => anyhow::bail!("invalid handle offset for var value: {voffset}"),
};
let mut size = std::mem::size_of::<String>()
+ std::mem::size_of::<Vec<u8>>()
+ key.len()
+ handle.length as usize;
for (k, v) in data.vars.iter() {
size += k.len();
size += v.len();
size += std::mem::size_of::<String>() + std::mem::size_of::<Vec<u8>>();
}
// If the store is larger than the configured size, or 1mb by default, then stop adding things
if size > data.manifest.memory.max_var_bytes.unwrap_or(1024 * 1024) as usize && voffset != 0 {
return Err(Error::msg("Variable store is full"));
}
let value = data.memory_bytes(handle)?.to_vec();
// Insert the value from memory into the `vars` map

View File

@@ -1,3 +1,5 @@
use extism_manifest::MemoryOptions;
use crate::*;
use std::{io::Write, time::Instant};
@@ -594,6 +596,29 @@ fn test_manifest_ptr_len() {
assert_eq!(count.get("count").unwrap().as_i64().unwrap(), 1);
}
#[test]
fn test_no_vars() {
let data = br#"
(module
(import "extism:host/env" "var_set" (func $var_set (param i64 i64)))
(import "extism:host/env" "input_offset" (func $input_offset (result i64)))
(func (export "test") (result i32)
(call $input_offset)
(call $input_offset)
(call $var_set)
(i32.const 0)
)
)
"#;
let manifest = Manifest::new([Wasm::data(data)])
.with_memory_options(MemoryOptions::new().with_max_var_bytes(1));
let mut plugin = Plugin::new(manifest, [], true).unwrap();
let output: Result<(), Error> = plugin.call("test", b"A".repeat(1024));
assert!(output.is_err());
let output: Result<(), Error> = plugin.call("test", vec![]);
assert!(output.is_ok());
}
#[test]
fn test_linking() {
let manifest = Manifest::new([