fix: get required bytes & grow memory before writing

Author: BowTiedWoo

clarity/src/vm/clarity_wasm.rs

@@ -549,6 +549,17 @@ pub fn call_function<'a>(
     }
     let mut in_mem_offset = offset + arg_size;
 
+    // Ensure that the memory has enough space for the arguments
+    let mut total_required_bytes = 0;
+    for (arg, ty) in args.iter().zip(func_types.get_arg_types()) {
+        total_required_bytes += get_required_bytes(ty, arg)?;
+    }
+    ensure_memory(
+        &memory,
+        &mut store,
+        total_required_bytes + in_mem_offset as usize,
+    )?;
+
     // Convert the args into wasmtime values
     let mut wasm_args = vec![];
     for (arg, ty) in args.iter().zip(func_types.get_arg_types()) {
@@ -1505,6 +1516,72 @@ fn ensure_memory(
     Ok(())
 }
 
+/// Get the number of bytes required to write the given value to memory.
+/// This is used to ensure that the memory has enough space for the arguments.
+fn get_required_bytes(ty: &TypeSignature, value: &Value) -> Result<usize, Error> {
+    match value {
+        Value::UInt(_) | Value::Int(_) | Value::Bool(_) => {
+            // These types don't require memory allocation
+            Ok(0)
+        }
+        Value::Optional(o) => {
+            let TypeSignature::OptionalType(inner_ty) = ty else {
+                return Err(Error::Wasm(WasmError::ValueTypeMismatch));
+            };
+
+            if let Some(inner_value) = o.data.as_ref() {
+                get_required_bytes(inner_ty, inner_value)
+            } else {
+                Ok(0)
+            }
+        }
+        Value::Response(r) => {
+            let TypeSignature::ResponseType(inner_tys) = ty else {
+                return Err(Error::Wasm(WasmError::ValueTypeMismatch));
+            };
+            get_required_bytes(
+                if r.committed {
+                    &inner_tys.0
+                } else {
+                    &inner_tys.1
+                },
+                &r.data,
+            )
+        }
+        Value::Sequence(SequenceData::String(CharType::ASCII(s))) => Ok(s.data.len()),
+        Value::Sequence(SequenceData::String(CharType::UTF8(s))) => Ok(s.data.len()),
+        Value::Sequence(SequenceData::Buffer(b)) => Ok(b.data.len()),
+        Value::Sequence(SequenceData::List(l)) => {
+            let TypeSignature::SequenceType(SequenceSubtype::ListType(ltd)) = ty else {
+                return Err(Error::Wasm(WasmError::ValueTypeMismatch));
+            };
+            let total_bytes = l
+                .data
+                .iter()
+                .map(|_| get_type_in_memory_size(ltd.get_list_item_type(), true))
+                .sum::<i32>() as usize;
+            Ok(total_bytes)
+        }
+        Value::Principal(PrincipalData::Standard(_)) => Ok(STANDARD_PRINCIPAL_BYTES),
+        Value::Principal(PrincipalData::Contract(p))
+        | Value::CallableContract(CallableData {
+            contract_identifier: p,
+            ..
+        }) => Ok(PRINCIPAL_BYTES + 1 + p.name.len() as usize),
+        Value::Tuple(TupleData { data_map, .. }) => {
+            let TypeSignature::TupleType(tuple_ty) = ty else {
+                return Err(Error::Wasm(WasmError::ValueTypeMismatch));
+            };
+
+            let mut total_bytes = 0;
+            for (name, ty) in tuple_ty.get_type_map() {
+                total_bytes += get_required_bytes(ty, &data_map[name])?;
+            }
+            Ok(total_bytes)
+        }
+    }
+}
+
 /// Convert a Clarity `Value` into one or more Wasm `Val`. If this value
 /// requires writing into the Wasm memory, write it to the provided `offset`.
 /// Return a vector of `Val`s that can be passed to a Wasm function, and the
@@ -1604,9 +1681,6 @@ fn pass_argument_to_wasm(
             Ok((buffer, new_offset, new_in_mem_offset))
         }
         Value::Sequence(SequenceData::String(CharType::ASCII(s))) => {
-            let required_bytes = (in_mem_offset as usize) + s.data.len();
-            ensure_memory(&memory, &mut store, required_bytes)?;
-
             // For a string, write the bytes into the memory, then pass the
             // offset and length to the Wasm function.
             let buffer = vec![Val::I32(in_mem_offset), Val::I32(s.data.len() as i32)];
@@ -1621,9 +1695,6 @@ fn pass_argument_to_wasm(
             Ok((buffer, offset, adjusted_in_mem_offset))
         }
         Value::Sequence(SequenceData::String(CharType::UTF8(s))) => {
-            let required_bytes = (in_mem_offset as usize) + s.data.len();
-            ensure_memory(&memory, &mut store, required_bytes)?;
-
             // For a utf8 string, convert the chars to big-endian i32, convert this into a list of
             // bytes, then pass the offset and length to the wasm function
             let bytes: Vec<u8> = String::from_utf8(s.items().iter().flatten().copied().collect())
@@ -1639,9 +1710,6 @@ fn pass_argument_to_wasm(
             Ok((buffer, offset, adjusted_in_mem_offset))
         }
         Value::Sequence(SequenceData::Buffer(b)) => {
-            let required_bytes = (in_mem_offset as usize) + b.data.len();
-            ensure_memory(&memory, &mut store, required_bytes)?;
-
             // For a buffer, write the bytes into the memory, then pass the
             // offset and length to the Wasm function.
             let buffer = vec![Val::I32(in_mem_offset), Val::I32(b.data.len() as i32)];
@@ -1659,14 +1727,6 @@ fn pass_argument_to_wasm(
             let TypeSignature::SequenceType(SequenceSubtype::ListType(ltd)) = ty else {
                 return Err(Error::Wasm(WasmError::ValueTypeMismatch));
             };
-            let total_bytes = l
-                .data
-                .iter()
-                .map(|_| get_type_in_memory_size(ltd.get_list_item_type(), true))
-                .sum::<i32>() as usize;
-            let required_bytes = (in_mem_offset as usize) + total_bytes;
-            ensure_memory(&memory, &mut store, required_bytes)?;
-
             let mut buffer = vec![Val::I32(offset)];
             let mut written = 0;
             let mut in_mem_written = 0;