Merge pull request #112 from dusk-network/fix-alloc

Fix `allocate` and `free_mem`

src/ffi.rs

@@ -6,7 +6,10 @@
 
 //! FFI bindings exposed to WASM module.
 
-use alloc::{vec, vec::Vec};
+use alloc::{
+    alloc::{alloc, dealloc, Layout},
+    vec::Vec,
+};
 use core::mem;
 
 use dusk_bytes::Serializable;
@@ -15,22 +18,28 @@ use sha2::{Digest, Sha512};
 
 use crate::{key, tx, types, utils, MAX_KEY, MAX_LEN};
 
+/// The alignment of the memory allocated by the FFI.
+///
+/// This is 1 because we're not allocating any complex data structures, and
+/// just interacting with the memory directly.
+const ALIGNMENT: usize = 1;
+
 /// Allocates a buffer of `len` bytes on the WASM memory.
 #[no_mangle]
 pub fn allocate(len: i32) -> i32 {
-    let bytes = vec![0u8; len as usize];
-    let ptr = bytes.as_ptr();
-    mem::forget(bytes);
-    ptr as i32
+    unsafe {
+        let layout = Layout::from_size_align_unchecked(len as usize, ALIGNMENT);
+        let ptr = alloc(layout);
+        ptr as _
+    }
 }
 
 /// Frees a previously allocated buffer on the WASM memory.
 #[no_mangle]
 pub fn free_mem(ptr: i32, len: i32) {
-    let ptr = ptr as *mut u8;
-    let len = len as usize;
     unsafe {
-        Vec::from_raw_parts(ptr, len, len);
+        let layout = Layout::from_size_align_unchecked(len as usize, ALIGNMENT);
+        dealloc(ptr as _, layout);
     }
 }
 
@@ -53,11 +62,9 @@ pub fn seed(args: i32, len: i32) -> i64 {
     hash.update(b"SEED");
 
     let seed = hash.finalize().to_vec();
-    let ptr = seed.as_ptr() as u32;
-    let len = seed.len() as u32;
 
-    mem::forget(seed);
-    utils::compose(true, ptr, len)
+    let (ptr, len) = utils::allocated_copy(seed);
+    utils::compose(true, ptr as _, len as _)
 }
 
 /// Computes the total balance of the given notes.

src/utils.rs

@@ -6,10 +6,10 @@
 
 //! Misc utilities required by the library implementation.
 
-use crate::{tx, MAX_INPUT_NOTES, MAX_LEN, RNG_SEED};
+use crate::{ffi, tx, MAX_INPUT_NOTES, MAX_LEN, RNG_SEED};
 
 use alloc::vec::Vec;
-use core::mem;
+use core::ptr;
 
 use dusk_bytes::DeserializableSlice;
 use dusk_jubjub::JubJubScalar;
@@ -93,14 +93,9 @@ pub fn into_ptr<T>(response: T) -> i64
 where
     T: Serialize,
 {
-    let response = serde_json::to_string(&response).unwrap_or_default();
-    let ptr = response.as_ptr() as u32;
-    let len = response.len() as u32;
-    let result = compose(true, ptr, len);
-
-    mem::forget(response);
-
-    result
+    let response = serde_json::to_string(&response).unwrap_or_default().leak();
+    let (ptr, len) = allocated_copy(response);
+    compose(true, ptr as _, len as _)
 }
 
 /// Returns the provided bytes as a pointer
@@ -113,13 +108,25 @@ where
         Err(_) => return fail(),
     };
 
-    let ptr = bytes.as_ptr() as u32;
-    let len = bytes.len() as u32;
+    let (ptr, len) = allocated_copy(bytes);
 
-    mem::forget(bytes);
     compose(true, ptr, len)
 }
 
+/// Allocated a new buffer, copies the provided bytes to it, and returns the
+/// pointer and length of the new buffer.
+pub fn allocated_copy<B: AsRef<[u8]>>(bytes: B) -> (u32, u32) {
+    unsafe {
+        let bytes = bytes.as_ref();
+        let len = bytes.len();
+
+        let ptr = ffi::allocate(bytes.len() as _);
+        ptr::copy_nonoverlapping(bytes.as_ptr(), ptr as _, len);
+
+        (ptr as _, len as _)
+    }
+}
+
 /// Creates a secure RNG directly a seed.
 pub fn rng(seed: [u8; 32]) -> ChaCha12Rng {
     ChaCha12Rng::from_seed(seed)