Skip to content

Commit

Permalink
wallet-core: enable WASM host to calc balance for a profile
Browse files Browse the repository at this point in the history
- Add `map_owned` FFI
- Add `balance` FFI
- Add `OwnedList` serializable type
- Add `ErrorCode` usage
- Move `map_owned` in `notes` module with `OwnedList`

Resolves #2317
  • Loading branch information
ZER0 committed Sep 9, 2024
1 parent 13b91c6 commit 202145b
Show file tree
Hide file tree
Showing 4 changed files with 215 additions and 56 deletions.
119 changes: 104 additions & 15 deletions wallet-core/src/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,40 @@
//
// Copyright (c) DUSK NETWORK. All rights reserved.

use crate::keys::{derive_bls_pk, derive_phoenix_pk, RNG_SEED};
use core::ptr;
//! This module provides the foreign function interface (FFI) for exposing
//! public functions from the `wallet-core` Rust library to a WASM runtime.
//! In addition to cryptographic operations, it offers memory management
//! functions, such as `malloc` and `free`, for interacting with the WASM
//! memory.
//!
//! This FFI allows seamless integration between Rust code and a WASM runtime
//! while ensuring efficient memory handling and secure key management.
#[macro_use]
pub(crate) mod debug;

pub mod error;
pub mod panic;

use crate::keys::{
derive_bls_pk, derive_phoenix_pk, derive_phoenix_sk, derive_phoenix_vk,
};
use crate::notes;
use crate::phoenix_balance;
use crate::Seed;
use error::ErrorCode;

use alloc::alloc::{alloc, dealloc, Layout};
use alloc::vec::Vec;
use core::{ptr, slice};
use dusk_bytes::Serializable;
use execution_core::{
signatures::bls::PublicKey as BlsPublicKey,
transfer::phoenix::PublicKey as PhoenixPublicKey,
transfer::phoenix::{NoteLeaf, PublicKey as PhoenixPublicKey},
};
use zeroize::Zeroize;

use alloc::alloc::{alloc, dealloc, Layout};
use rkyv::{from_bytes, to_bytes};

/// The alignment of the memory allocated by the FFI.
///
Expand All @@ -39,12 +64,33 @@ pub fn free(ptr: u32, len: u32) {
}
}

/// Map a list of indexes into keys using the provided seed and callback.
unsafe fn indexes_into_keys<T, F>(
seed: &Seed,
indexes: *const u8,
mut callback: F,
) -> Vec<T>
where
F: FnMut(&Seed, u8) -> T,
{
let len = *indexes as usize;
let slice = slice::from_raw_parts(indexes.add(1), len);
slice.iter().map(|&byte| callback(seed, byte)).collect()
}

unsafe fn read_buffer(ptr: *const u8) -> Vec<u8> {
let len = slice::from_raw_parts(ptr, 4);
let len = u32::from_le_bytes(len.try_into().unwrap()) as usize;
slice::from_raw_parts(ptr.add(4), len).to_vec()
}

/// Generate a profile (account / address pair) for the given seed and index.
#[no_mangle]
pub unsafe extern "C" fn generate_profile(
seed: &[u8; RNG_SEED],
seed: &Seed,
index: u8,
profile: *mut [u8; PhoenixPublicKey::SIZE + BlsPublicKey::SIZE],
) -> u8 {
) -> ErrorCode {
let ppk = derive_phoenix_pk(seed, index).to_bytes();
let bpk = derive_bls_pk(seed, index).to_bytes();

Expand All @@ -60,16 +106,59 @@ pub unsafe extern "C" fn generate_profile(
BlsPublicKey::SIZE,
);

0
ErrorCode::Ok
}

// Currently we're not handling panic message in the WASM module; in the future
// we might want to enable it for `debug` releases.
mod panic_handling {
use core::panic::PanicInfo;
/// Filter all notes and their block height that are owned by the given keys,
/// mapped to their nullifiers.
#[no_mangle]
pub unsafe fn map_owned(
seed: &Seed,
indexes: *const u8,
notes_ptr: *mut u8,
) -> ErrorCode {
let keys = indexes_into_keys(seed, indexes, derive_phoenix_sk);
let notes = read_buffer(notes_ptr);
let notes: Vec<NoteLeaf> = from_bytes::<Vec<NoteLeaf>>(&notes)
.or(Err(ErrorCode::UnarchivingError))?;

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
loop {}
}
let owned = notes::map_owned(&keys, notes);

keys.into_iter().for_each(|mut sk| sk.zeroize());

let bytes =
to_bytes::<_, 4096>(&owned).or(Err(ErrorCode::ArchivingError))?;

let len = bytes.len().to_le_bytes();

ptr::copy_nonoverlapping(len.as_ptr(), notes_ptr, 4);
ptr::copy_nonoverlapping(bytes.as_ptr(), notes_ptr.add(4), bytes.len());

ErrorCode::Ok
}

/// Calculate the balance info for the phoenix address at the given index for
/// the given seed.
#[no_mangle]
pub unsafe fn balance(
seed: &Seed,
index: u8,
notes_ptr: *const u8,
balance_info_ptr: *mut [u8; 16],
) -> ErrorCode {
let vk = derive_phoenix_vk(seed, index);

let notes = read_buffer(notes_ptr);
let notes: Vec<NoteLeaf> = from_bytes::<Vec<NoteLeaf>>(&notes)
.or(Err(ErrorCode::UnarchivingError))?;

let info = phoenix_balance(&vk, notes.iter());

ptr::copy_nonoverlapping(
info.to_bytes().as_ptr(),
&mut (*balance_info_ptr)[0],
16,
);

ErrorCode::Ok
}
22 changes: 9 additions & 13 deletions wallet-core/src/keys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,17 +21,13 @@ use execution_core::{
},
};

/// The seed bytes buffer which is used at multiple places
pub type Seed = [u8; RNG_SEED];

/// Length of the seed of the generated rng.
pub const RNG_SEED: usize = 64;
use crate::Seed;

/// Generates a [`BlsSecretKey`] from a seed and index.
///
/// The randomness is generated using [`rng_with_index`].
#[must_use]
pub fn derive_bls_sk(seed: &[u8; RNG_SEED], index: u8) -> BlsSecretKey {
pub fn derive_bls_sk(seed: &Seed, index: u8) -> BlsSecretKey {
// note that if we change the string used for the rng, all previously
// generated keys will become invalid
// NOTE: When breaking the keys, we will want to change the string too
Expand All @@ -42,7 +38,7 @@ pub fn derive_bls_sk(seed: &[u8; RNG_SEED], index: u8) -> BlsSecretKey {
///
/// The randomness is generated using [`rng_with_index`].
#[must_use]
pub fn derive_bls_pk(seed: &[u8; RNG_SEED], index: u8) -> BlsPublicKey {
pub fn derive_bls_pk(seed: &Seed, index: u8) -> BlsPublicKey {
let mut sk = derive_bls_sk(seed, index);
let pk = BlsPublicKey::from(&sk);
sk.zeroize();
Expand All @@ -54,7 +50,7 @@ pub fn derive_bls_pk(seed: &[u8; RNG_SEED], index: u8) -> BlsPublicKey {
///
/// The randomness is generated using [`rng_with_index`].
#[must_use]
pub fn derive_phoenix_sk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixSecretKey {
pub fn derive_phoenix_sk(seed: &Seed, index: u8) -> PhoenixSecretKey {
// note that if we change the string used for the rng, all previously
// generated keys will become invalid
// NOTE: When breaking the keys, we will want to change the string too
Expand All @@ -66,21 +62,21 @@ pub fn derive_phoenix_sk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixSecretKey {
/// The randomness is generated using [`rng_with_index`].
#[must_use]
pub fn derive_multiple_phoenix_sk(
seed: &[u8; RNG_SEED],
seed: &Seed,
index_range: Range<u8>,
) -> Vec<PhoenixSecretKey> {
index_range
.map(|index| derive_phoenix_sk(seed, index))
.collect()
}

/// Generates a [`PheonixPublicKey`] from its seed and index.
/// Generates a [`PhoenixPublicKey`] from its seed and index.
///
/// First the [`PhoenixSecretKey`] is derived with [`derive_phoenix_sk`], then
/// the public key is generated from it and the secret key is erased from
/// memory.
#[must_use]
pub fn derive_phoenix_pk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixPublicKey {
pub fn derive_phoenix_pk(seed: &Seed, index: u8) -> PhoenixPublicKey {
let mut sk = derive_phoenix_sk(seed, index);
let pk = PhoenixPublicKey::from(&sk);
sk.zeroize();
Expand All @@ -93,7 +89,7 @@ pub fn derive_phoenix_pk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixPublicKey {
/// First the [`PhoenixSecretKey`] is derived with [`derive_phoenix_sk`], then
/// the view key is generated from it and the secret key is erased from memory.
#[must_use]
pub fn derive_phoenix_vk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixViewKey {
pub fn derive_phoenix_vk(seed: &Seed, index: u8) -> PhoenixViewKey {
let mut sk = derive_phoenix_sk(seed, index);
let vk = PhoenixViewKey::from(&sk);
sk.zeroize();
Expand All @@ -110,7 +106,7 @@ pub fn derive_phoenix_vk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixViewKey {
/// subsequently used to generate the key.
#[must_use]
pub fn rng_with_index(
seed: &[u8; RNG_SEED],
seed: &Seed,
index: u8,
termination: &[u8],
) -> ChaCha12Rng {
Expand Down
38 changes: 10 additions & 28 deletions wallet-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,57 +10,39 @@
#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
#![deny(clippy::pedantic)]
#![feature(try_trait_v2)]

#[cfg(target_family = "wasm")]
#[global_allocator]
static ALLOC: dlmalloc::GlobalDlmalloc = dlmalloc::GlobalDlmalloc;

extern crate alloc;

#[cfg(target_family = "wasm")]
#[macro_use]
mod ffi;

pub mod input;
pub mod keys;
pub mod notes;
pub mod transaction;

/// The seed used to generate the entropy for the keys
pub type Seed = [u8; 64];

pub mod prelude {
//! Re-export of the most commonly used types and traits.
pub use crate::keys;
pub use crate::{input::MAX_INPUT_NOTES, keys::RNG_SEED};
}

use alloc::collections::btree_map::BTreeMap;
use alloc::vec::Vec;

use dusk_bytes::{DeserializableSlice, Serializable, Write};

use execution_core::{
transfer::phoenix::{
Note, NoteLeaf, SecretKey as PhoenixSecretKey,
ViewKey as PhoenixViewKey,
},
BlsScalar,
};

/// Filter all notes and their block height that are owned by the given keys,
/// mapped to their nullifiers.
pub fn map_owned(
keys: impl AsRef<[PhoenixSecretKey]>,
notes: impl AsRef<[NoteLeaf]>,
) -> BTreeMap<BlsScalar, NoteLeaf> {
notes
.as_ref()
.iter()
.fold(BTreeMap::new(), |mut notes_map, note_leaf| {
for sk in keys.as_ref() {
if sk.owns(note_leaf.note.stealth_address()) {
let nullifier = note_leaf.note.gen_nullifier(sk);
notes_map.insert(nullifier, note_leaf.clone());
}
}
notes_map
})
}
use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey};

pub use notes::map_owned;

/// Calculate the sum for all the given [`Note`]s that belong to the given
/// [`PhoenixViewKey`].
Expand Down
92 changes: 92 additions & 0 deletions wallet-core/src/notes.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
// 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 http://mozilla.org/MPL/2.0/.
//
// Copyright (c) DUSK NETWORK. All rights reserved.

//! Implementations of basic wallet functionalities to create transactions.
use alloc::vec::Vec;
use core::ops::Index;
use execution_core::{
transfer::phoenix::{NoteLeaf, SecretKey as PhoenixSecretKey},
BlsScalar,
};

use rkyv::{Archive, Deserialize, Serialize};

/// A collection of notes stored as key-value pairs.
/// The key is a `BlsScalar` and the value is a `NoteLeaf`.
/// Duplicates are allowed.
#[derive(Default, Archive, Serialize, Deserialize, Debug)]
pub struct OwnedList {
/// The underlying storage of key-value pairs where
/// `BlsScalar` is the key and `NoteLeaf` is the value.
entries: Vec<(BlsScalar, NoteLeaf)>,
}

impl OwnedList {
/// Inserts a new key-value pair into the collection.
pub fn insert(&mut self, key: BlsScalar, value: NoteLeaf) {
self.entries.push((key, value));
}

/// Returns the number of entries (key-value pairs) in the collection.
#[must_use]
pub fn len(&self) -> usize {
self.entries.len()
}

/// Checks if the collection is empty.
#[must_use]
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}

/// Retrieves the value (`NoteLeaf`) associated with a given key
#[must_use]
pub fn get(&self, key: &BlsScalar) -> Option<&NoteLeaf> {
self.entries.iter().find(|(k, _)| k == key).map(|(_, v)| v)
}

/// Retrieves all keys in the collection.
#[must_use]
pub fn keys(&self) -> Vec<BlsScalar> {
self.entries.iter().map(|(k, _)| *k).collect()
}
}

impl Index<&BlsScalar> for OwnedList {
type Output = NoteLeaf;

/// Retrieves the value (`NoteLeaf`) associated with a given key
/// (`BlsScalar`).
///
/// Panics if the key is not found in the collection.
fn index(&self, index: &BlsScalar) -> &Self::Output {
self.get(index).expect("key not found")
}
}

/// Filter all notes and their block height that are owned by the given keys,
/// mapped to their nullifiers.
pub fn map_owned(
keys: impl AsRef<[PhoenixSecretKey]>,
notes: impl AsRef<[NoteLeaf]>,
) -> OwnedList {
notes.as_ref().iter().fold(
OwnedList::default(),
|mut notes_map, note_leaf| {
eprintln!("Printing note...");
dbg!(note_leaf);
for sk in keys.as_ref() {
if sk.owns(note_leaf.note.stealth_address()) {
let nullifier = note_leaf.note.gen_nullifier(sk);
notes_map.insert(nullifier, note_leaf.clone());
break;
}
}
notes_map
},
)
}

0 comments on commit 202145b

Please sign in to comment.