Add Safety comment, set MAX_KEY to range exclusive

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add support for phoenix_core methods
 - Add support for stake contract functions
 - Add support for transfer contract functions
+- Update dependencies for new phoenix_core
 
 ### Changed
 - Change the whole api to support `wasm32-unkonwn-unknown`
@@ -22,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix the input selecting algorithm from custom new one to old wallet-core to fix wrong picking of notes
 - Fix the execute to work with only one public spend key and one `Transaction`
 - Fix double rkyv serialization bug #113
+- Fix memory allocation bug with `mem::forget`
 
 ## [Old Changelog below]
 
@@ -133,6 +135,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Expose `NodeClient` and `Store` through FFI
 - Define FFI and compile it only for WASM
 
+[#113]: https://github.com/dusk-network/wallet-core/issues/113
 [#58]: https://github.com/dusk-network/wallet-core/issues/58
 [#55]: https://github.com/dusk-network/wallet-core/issues/55
 [#53]: https://github.com/dusk-network/wallet-core/issues/53
@@ -141,7 +144,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#44]: https://github.com/dusk-network/wallet-core/issues/44
 [#41]: https://github.com/dusk-network/wallet-core/issues/41
 [#40]: https://github.com/dusk-network/wallet-core/issues/40
-[#113]: https://github.com/dusk-network/wallet-core/issues/113
 [#34]: https://github.com/dusk-network/wallet-core/issues/34
 [#31]: https://github.com/dusk-network/wallet-core/issues/31
 [#25]: https://github.com/dusk-network/wallet-core/issues/25

src/compat/crypto.rs

@@ -29,14 +29,16 @@ const TREE_LEAF_SIZE: usize = size_of::<ArchivedTreeLeaf>();
 /// if its true then nullifier of that note if sent with it
 #[no_mangle]
 pub fn check_note_ownership(args: i32, len: i32) -> i64 {
-    let args = utils::take_args_raw(args, len);
+    // SAFETY: We assume the caller has passed a valid pointer and len as the
+    // function arguments else we might get undefined behavior
+    let args = unsafe { core::slice::from_raw_parts(args as _, len as _) };
 
     let seed = &args[..64];
     let leaves: &[u8] = &args[64..];
 
-    let seed = match seed.try_into().ok() {
-        Some(s) => s,
-        None => return utils::fail(),
+    let seed = match seed.try_into() {
+        Ok(s) => s,
+        Err(_) => return utils::fail(),
     };
 
     let mut leaf_chunk = leaves.chunks_exact(TREE_LEAF_SIZE);
@@ -46,24 +48,31 @@ pub fn check_note_ownership(args: i32, len: i32) -> i64 {
     let mut nullifiers = Vec::new();
     let mut block_heights = Vec::new();
     let mut public_spend_keys = Vec::new();
+    let mut view_keys = Vec::with_capacity(MAX_KEY);
+    let mut secret_keys = Vec::with_capacity(MAX_KEY);
+
+    for idx in 0..MAX_KEY {
+        let idx = idx as u64;
+        let view_key = key::derive_vk(&seed, idx);
+        let sk = key::derive_sk(&seed, idx as _);
+        view_keys.push(view_key);
+        secret_keys.push(sk);
+    }
 
     for leaf_bytes in leaf_chunk.by_ref() {
-        let TreeLeaf { block_height, note } =
-            match rkyv::from_bytes(leaf_bytes).ok() {
-                Some(a) => a,
-                None => {
-                    return utils::fail();
-                }
-            };
+        let TreeLeaf { block_height, note } = match rkyv::from_bytes(leaf_bytes)
+        {
+            Ok(a) => a,
+            Err(_) => {
+                return utils::fail();
+            }
+        };
 
         last_pos = core::cmp::max(last_pos, *note.pos());
 
         for idx in 0..=MAX_KEY {
-            let idx = idx as u64;
-            let view_key = key::derive_vk(&seed, idx);
-
-            if view_key.owns(&note) {
-                let sk = key::derive_sk(&seed, idx);
+            if view_keys[idx].owns(&note) {
+                let sk = secret_keys[idx];
                 let nullifier = note.gen_nullifier(&sk);
 
                 let nullifier_found =
@@ -77,9 +86,9 @@ pub fn check_note_ownership(args: i32, len: i32) -> i64 {
                     bs58::encode(PublicKey::from(sk).to_bytes()).into_string();
 
                 let raw_note: Vec<u8> =
-                    match rkyv::to_bytes::<Note, MAX_LEN>(&note).ok() {
-                        Some(n) => n.to_vec(),
-                        None => return utils::fail(),
+                    match rkyv::to_bytes::<Note, MAX_LEN>(&note) {
+                        Ok(n) => n.to_vec(),
+                        Err(_) => return utils::fail(),
                     };
 
                 notes.push(raw_note.to_owned());

src/ffi.rs

@@ -92,15 +92,15 @@ pub fn balance(args: i32, len: i32) -> i64 {
         Err(_) => return utils::fail(),
     };
 
-    let mut keys = unsafe { [mem::zeroed(); MAX_KEY + 1] };
+    let mut keys = unsafe { [mem::zeroed(); MAX_KEY] };
     let mut values = Vec::with_capacity(notes.len());
     let mut keys_len = 0;
     let mut sum = 0u64;
 
     'outer: for note in notes {
         // we iterate all the available keys until one can successfully decrypt
         // the note. if all fails, returns false
-        for idx in 0..=MAX_KEY {
+        for idx in 0..MAX_KEY {
             if keys_len == idx {
                 keys[idx] = key::derive_vk(&seed, idx as u64);
                 keys_len += 1;
@@ -355,7 +355,7 @@ pub fn public_keys(args: i32, len: i32) -> i64 {
         None => return utils::fail(),
     };
 
-    let keys = (0..=MAX_KEY)
+    let keys = (0..MAX_KEY)
         .map(|idx| key::derive_pk(&seed, idx as u64))
         .map(|pk| bs58::encode(pk.to_bytes()).into_string())
         .collect();
@@ -382,7 +382,7 @@ pub fn view_keys(args: i32, len: i32) -> i64 {
         None => return utils::fail(),
     };
 
-    let keys: Vec<_> = (0..=MAX_KEY)
+    let keys: Vec<_> = (0..MAX_KEY)
         .map(|idx| key::derive_vk(&seed, idx as u64))
         .collect();
 
@@ -416,14 +416,14 @@ pub fn nullifiers(args: i32, len: i32) -> i64 {
     };
 
     let mut nullifiers = Vec::with_capacity(notes.len());
-    let mut sks = unsafe { [mem::zeroed(); MAX_KEY + 1] };
-    let mut vks = unsafe { [mem::zeroed(); MAX_KEY + 1] };
+    let mut sks = unsafe { [mem::zeroed(); MAX_KEY] };
+    let mut vks = unsafe { [mem::zeroed(); MAX_KEY] };
     let mut keys_len = 0;
 
     'outer: for note in notes {
         // we iterate all the available view key until one can successfully
         // decrypt the note. if any fails, returns false
-        for idx in 0..=MAX_KEY {
+        for idx in 0..MAX_KEY {
             if keys_len == idx {
                 sks[idx] = key::derive_sk(&seed, idx as u64);
                 vks[idx] = ViewKey::from(&sks[idx]);

src/utils.rs

@@ -63,14 +63,6 @@ where
     serde_json::from_str(&args).ok()
 }
 
-/// reads the raw bytes at the pointer for the length and returns what it reason
-pub fn take_args_raw<'a>(args: i32, len: i32) -> &'a [u8] {
-    let args = args as *mut u8;
-    let len = len as usize;
-
-    unsafe { core::slice::from_raw_parts(args, len) }
-}
-
 /// Sanitizes arbitrary bytes into well-formed seed.
 pub fn sanitize_seed(bytes: Vec<u8>) -> Option<[u8; RNG_SEED]> {
     (bytes.len() == RNG_SEED).then(|| {

tests/wallet.rs

@@ -200,7 +200,7 @@ fn public_keys_works() {
         PublicKey::from_bytes(&key_array).unwrap();
     }
 
-    assert_eq!(keys.len(), MAX_KEY + 1);
+    assert_eq!(keys.len(), MAX_KEY);
 }
 
 #[test]
@@ -336,8 +336,7 @@ mod node {
             .into_iter()
             .map(|value| {
                 let obfuscated = (rng.next_u32() & 1) == 1;
-                let idx = rng.next_u64()
-                    % (if MAX_KEY == 0 { 1 } else { MAX_KEY }) as u64;
+                let idx = rng.next_u64() % (MAX_KEY) as u64;
                 let sk = key::derive_sk(seed, idx);
                 let pk = PublicKey::from(&sk);