From d47af2455208ddd0923bf4d0c02e9eef6f92622c Mon Sep 17 00:00:00 2001
From: zach <zach@dylibso.com>
Date: Thu, 7 Mar 2024 09:55:02 -0800
Subject: [PATCH] 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
---
 manifest/schema.json         | 13 ++++++++++++-
 manifest/src/lib.rs          | 34 ++++++++++++++++++++++++++++++++++
 runtime/src/pdk.rs           | 28 +++++++++++++++++++---------
 runtime/src/tests/runtime.rs | 25 +++++++++++++++++++++++++
 4 files changed, 90 insertions(+), 10 deletions(-)

diff --git a/manifest/schema.json b/manifest/schema.json
index c8d5513..c988a24 100644
--- a/manifest/schema.json
+++ b/manifest/schema.json
@@ -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
diff --git a/manifest/src/lib.rs b/manifest/src/lib.rs
index 85bcfbb..a53ef89 100644
--- a/manifest/src/lib.rs
+++ b/manifest/src/lib.rs
@@ -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
diff --git a/runtime/src/pdk.rs b/runtime/src/pdk.rs
index 3816d11..5e10e5c 100644
--- a/runtime/src/pdk.rs
+++ b/runtime/src/pdk.rs
@@ -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
diff --git a/runtime/src/tests/runtime.rs b/runtime/src/tests/runtime.rs
index bc1b1d1..42b86d2 100644
--- a/runtime/src/tests/runtime.rs
+++ b/runtime/src/tests/runtime.rs
@@ -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([