fix: Handle inputs of moderate size and up (#91)

Handles inputs of moderate size and up by ensuring that the input byte array is fully cloned before parsing it as json.

Author: AndrewSisley

.gitignore

@@ -19,6 +19,9 @@ tests/modules/rust_wasm32_filter/pkg
 tests/modules/rust_wasm32_normalize/Cargo.lock
 tests/modules/rust_wasm32_normalize/target
 tests/modules/rust_wasm32_normalize/pkg
+tests/modules/rust_wasm32_memory/Cargo.lock
+tests/modules/rust_wasm32_memory/target
+tests/modules/rust_wasm32_memory/pkg
 sdk-rust/target
 sdk-rust/Cargo.lock
 host-go/build/host-go.exe

sdk-rust/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "lens_sdk"
 description = "An SDK for writing bi-directional, wasm based, LensVM transformations in Rust"
-version = "0.6.0"
+version = "0.7.0"
 edition = "2021"
 readme = "../README.md"
 repository = "https://github.com/lens-vm/lens"
@@ -12,5 +12,5 @@ crate-type = ["lib"]
 
 [dependencies]
 serde = "1.0"
-serde_json = "1.0.87"
-byteorder = "1.4.3"
+serde_json = "1.0"
+byteorder = "1.4"

sdk-rust/src/lib.rs

@@ -82,7 +82,7 @@ pub fn free(ptr: *mut u8, size: usize) {
         Vec::from_raw_parts(ptr, size, size)
     };
     let buf = ManuallyDrop::new(buf);
-    mem::drop(ManuallyDrop::into_inner(buf));
+    ManuallyDrop::into_inner(buf);
 }
 
 /// Manually drop the memory occupied by a transport buffer at the given location.
@@ -116,15 +116,15 @@ pub fn free_transport_buffer(ptr: *mut u8) -> Result<()> {
 
     let type_id: i8  = type_rdr.read_i8()?;
     if type_id == NIL_TYPE_ID {
-        mem::drop(ManuallyDrop::into_inner(type_rdr));
+        ManuallyDrop::into_inner(type_rdr);
         return Ok(())
     }
     if type_id == EOS_TYPE_ID {
-        mem::drop(ManuallyDrop::into_inner(type_rdr));
+        ManuallyDrop::into_inner(type_rdr);
         return Ok(())
     }
     if type_id < 0 {
-        mem::drop(ManuallyDrop::into_inner(type_rdr));
+        ManuallyDrop::into_inner(type_rdr);
         return Ok(())
     }
 
@@ -137,7 +137,7 @@ pub fn free_transport_buffer(ptr: *mut u8) -> Result<()> {
 
     let len: usize = len_rdr.read_u32::<LittleEndian>()?.try_into()?;
 
-    free(ptr, len);
+    free(ptr, mem::size_of::<i8>()+mem::size_of::<u32>()+len);
 
     Ok(())
 }
@@ -162,51 +162,49 @@ pub fn free_transport_buffer(ptr: *mut u8) -> Result<()> {
 ///
 /// This function will return an [Error](error/enum.Error.html) if the data at the given location is not in the expected
 /// format.
-pub fn try_from_mem<TOutput: for<'a> Deserialize<'a>>(ptr: *mut u8) -> Result<StreamOption<TOutput>> {
+pub fn try_from_mem<TOutput: for<'a> Deserialize<'a> + Clone>(ptr: *mut u8) -> Result<StreamOption<TOutput>> {
     let type_vec: Vec<u8> = unsafe {
         Vec::from_raw_parts(ptr, mem::size_of::<i8>(), mem::size_of::<i8>())
     };
 
-    let type_rdr = Cursor::new(type_vec);
-    let mut type_rdr = ManuallyDrop::new(type_rdr);
+    let mut type_rdr = Cursor::new(type_vec.clone());
+    let _ = ManuallyDrop::new(type_vec);
 
     let type_id: i8  = type_rdr.read_i8()?;
     if type_id == NIL_TYPE_ID {
-        mem::drop(ManuallyDrop::into_inner(type_rdr));
+        free_transport_buffer(ptr)?;
         return Ok(StreamOption::None)
     }
     if type_id == EOS_TYPE_ID {
-        mem::drop(ManuallyDrop::into_inner(type_rdr));
+        free_transport_buffer(ptr)?;
         return Ok(StreamOption::EndOfStream)
     }
     if type_id < 0 {
-        mem::drop(ManuallyDrop::into_inner(type_rdr));
+        free_transport_buffer(ptr)?;
         return Result::from(error::LensError::InputErrorUnsupportedError)
     }
 
     let len_vec: Vec<u8> = unsafe {
         Vec::from_raw_parts(ptr.add(mem::size_of::<i8>()), mem::size_of::<u32>(), mem::size_of::<u32>())
     };
 
-    let len_rdr = Cursor::new(len_vec);
-    let mut len_rdr = ManuallyDrop::new(len_rdr);
+    let mut len_rdr = Cursor::new(len_vec.clone());
+    let _ = ManuallyDrop::new(len_vec);
 
     let len: usize = len_rdr.read_u32::<LittleEndian>()?.try_into()?;
 
     let input_vec: Vec<u8> = unsafe {
         Vec::from_raw_parts(ptr.add(mem::size_of::<i8>()+mem::size_of::<u32>()), len, len)
     };
-    let input_vec = ManuallyDrop::new(input_vec);
-
-    let json_string = String::from_utf8(input_vec.to_vec())?.clone();
 
     // It is possible for null json values to reach this line, particularly if sourced directly
     // from a 3rd party module, so we ensure that we parse to option as well as the earlier type_id
     // checks.
-    let result = match serde_json::from_str::<Option<TOutput>>(&json_string)? {
-        Some(v) => StreamOption::Some(v),
+    let result = match serde_json::from_slice::<Option<TOutput>>(&input_vec.clone())? {
+        Some(v) => StreamOption::Some(v.clone()),
         None => StreamOption::None,
     };
+    let _ = ManuallyDrop::new(input_vec);
 
     Ok(result)
 }
@@ -236,7 +234,8 @@ pub fn try_to_mem(type_id: i8, message: &[u8]) -> Result<*mut u8> {
         }
     };
 
-    let result = wtr.into_inner().clone().as_mut_ptr();
+    let mut buffer = wtr.into_inner().clone();
+    let result = buffer.as_mut_ptr();
     Ok(result)
 }
 

tests/integration/cli/with_mem_test.go

@@ -0,0 +1,47 @@
+package integration
+
+import (
+	"testing"
+
+	"github.com/lens-vm/lens/tests/modules"
+)
+
+// Rust is very aggressive in cleaning up memory, and we had a bug where lenses would fail if the input
+// was particularly large, and/or the transform copied the input too many times.
+func TestWithMem(t *testing.T) {
+	type Value struct {
+		Name string
+		Type string `json:"__type"`
+		// Inline arrays seem to be particularly problematic, I guess it is due to some lifetime stuff
+		// in the serde library.
+		Array []string
+	}
+
+	executeTest(
+		t,
+		TestCase[Value, Value]{
+			LensFile: `
+			{
+				"lenses": [
+					{
+						"path": "` + modules.WasmPath_Memory + `"
+					}
+				]
+			}`,
+			Input: []Value{
+				{
+					Name:  "John",
+					Type:  "passsssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+					Array: []string{},
+				},
+			},
+			ExpectedOutput: []Value{
+				{
+					Name:  "John",
+					Type:  "passsssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss",
+					Array: []string{},
+				},
+			},
+		},
+	)
+}

tests/modules/Makefile

@@ -11,4 +11,5 @@ build:
 	cargo build --target wasm32-unknown-unknown --manifest-path "./rust_wasm32_counter/Cargo.toml"
 	cargo build --target wasm32-unknown-unknown --manifest-path "./rust_wasm32_filter/Cargo.toml"
 	cargo build --target wasm32-unknown-unknown --manifest-path "./rust_wasm32_normalize/Cargo.toml"
+	cargo build --target wasm32-unknown-unknown --manifest-path "./rust_wasm32_memory/Cargo.toml"
 	(cd "./as_wasm32_simple/" && npm install && npm run asbuild:debug)

tests/modules/rust_wasm32_counter/src/lib.rs

@@ -10,7 +10,7 @@ extern "C" {
     fn next() -> *mut u8;
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct Value {
     #[serde(rename = "Id")]
 	pub id: usize,

tests/modules/rust_wasm32_filter/src/lib.rs

@@ -8,7 +8,7 @@ extern "C" {
     fn next() -> *mut u8;
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct Value {
     #[serde(rename = "Name")]
     pub name: String,

tests/modules/rust_wasm32_memory/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "rust-wasm32-memory"
+version = "0.1.0"
+edition = "2018"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+lens_sdk = { path = "../../../sdk-rust" }

tests/modules/rust_wasm32_memory/src/lib.rs

@@ -0,0 +1,61 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+use std::error::Error;
+use serde::{Deserialize, Serialize};
+use lens_sdk::StreamOption;
+use lens_sdk::option::StreamOption::{Some, None, EndOfStream};
+
+#[link(wasm_import_module = "lens")]
+extern "C" {
+    fn next() -> *mut u8;
+}
+
+#[no_mangle]
+pub extern "C" fn alloc(size: usize) -> *mut u8 {
+    lens_sdk::alloc(size)
+}
+
+#[no_mangle]
+pub extern "C" fn transform() -> *mut u8 {
+    match try_transform() {
+        Ok(o) => match o {
+            Some(result_json) => lens_sdk::to_mem(lens_sdk::JSON_TYPE_ID, &result_json),
+            None => lens_sdk::nil_ptr(),
+            EndOfStream => lens_sdk::to_mem(lens_sdk::EOS_TYPE_ID, &[]),
+        },
+        Err(e) => lens_sdk::to_mem(lens_sdk::ERROR_TYPE_ID, &e.to_string().as_bytes())
+    }
+}
+
+#[derive(Serialize, Deserialize, Clone)]
+pub struct Input {
+    #[serde(rename = "Name")]
+    pub name: String,
+    #[serde(rename= "__type")]
+	pub __type: String,
+    #[serde(rename = "Array")]
+	pub array: Vec<String>,
+}
+
+fn try_transform() -> Result<StreamOption<Vec<u8>>, Box<dyn Error>> {
+    let ptr = unsafe { next() };
+    let mut input = match lens_sdk::try_from_mem::<Input>(ptr)? {
+        Some(v) => v,
+        // Implementations of `transform` are free to handle nil however they like. In this
+        // implementation we chose to return nil given a nil input.
+        None => return Ok(None),
+        EndOfStream => return Ok(EndOfStream)
+    };
+
+    for _ in 1..1000 {
+        // Copy the input a bunch of times to make sure that the memory is not cleaned up
+        // before we are finished with it.
+        input = input.clone();
+    }
+
+    let result_json = serde_json::to_vec(&input.clone())?;
+    lens_sdk::free_transport_buffer(ptr)?;
+    Ok(Some(result_json))
+}

tests/modules/rust_wasm32_normalize/src/lib.rs

@@ -11,7 +11,7 @@ extern "C" {
     fn next() -> *mut u8;
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct Book {
     #[serde(rename = "Name")]
     pub name: String,