Skip to content

Commit

Permalink
wallet-core: Add filter_map function
Browse files Browse the repository at this point in the history
Resolves #2129
  • Loading branch information
moCello committed Aug 14, 2024
1 parent 0bb889e commit ba594e6
Show file tree
Hide file tree
Showing 5 changed files with 217 additions and 75 deletions.
18 changes: 17 additions & 1 deletion wallet-core/src/keys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,12 @@

//! Utilities to derive keys from the seed.
use alloc::vec::Vec;
use core::ops::Range;

use rand_chacha::{rand_core::SeedableRng, ChaCha12Rng};
use sha2::{Digest, Sha256};
use zeroize::Zeroize;

use execution_core::{
signatures::bls::SecretKey as BlsSecretKey,
Expand All @@ -16,7 +20,6 @@ use execution_core::{
ViewKey as PhoenixViewKey,
},
};
use zeroize::Zeroize;

use crate::RNG_SEED;

Expand All @@ -40,6 +43,19 @@ pub fn derive_phoenix_sk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixSecretKey {
PhoenixSecretKey::random(&mut rng_with_index(seed, index, b"PSK"))
}

/// Generates multiple [`PhoenixSecretKey`] from a seed and a range of indices.
///
/// The randomness is generated using [`rng_with_index`].
#[must_use]
pub fn derive_multiple_phoenix_sk(
seed: &[u8; RNG_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.
///
/// First the [`PhoenixSecretKey`] is derived with [`derive_phoenix_sk`], then
Expand Down
28 changes: 27 additions & 1 deletion wallet-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,35 @@ pub const RNG_SEED: usize = 64;
// phoenix-transaction
const MAX_INPUT_NOTES: usize = 4;

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

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

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

/// Calculate the sum for all the given [`Note`]s that belong to the given
/// [`PhoenixViewKey`].
Expand Down
72 changes: 0 additions & 72 deletions wallet-core/tests/balance.rs

This file was deleted.

26 changes: 25 additions & 1 deletion wallet-core/tests/keys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
use dusk_bytes::Serializable;

use wallet_core::keys::{
derive_bls_sk, derive_phoenix_pk, derive_phoenix_sk, derive_phoenix_vk,
derive_bls_sk, derive_multiple_phoenix_sk, derive_phoenix_pk,
derive_phoenix_sk, derive_phoenix_vk,
};

const SEED: [u8; 64] = [0; 64];
Expand All @@ -26,6 +27,29 @@ fn test_derive_phoenix_sk() {
assert_eq!(derive_phoenix_sk(&SEED, INDEX).to_bytes(), sk_bytes);
}

#[test]
fn test_derive_multiple_phoenix_sk() {
// it is important that we always derive the same key from a fixed seed
let sk_bytes_0 = [
160, 210, 234, 8, 94, 23, 76, 60, 130, 143, 137, 225, 37, 83, 68, 218,
207, 192, 171, 235, 252, 130, 133, 62, 18, 232, 6, 49, 245, 123, 220,
12, 250, 111, 39, 88, 24, 41, 156, 174, 241, 14, 118, 173, 11, 53, 192,
126, 7, 119, 70, 69, 212, 230, 124, 79, 223, 140, 93, 153, 33, 147,
163, 0,
];
let sk_bytes_1 = [
0, 229, 97, 222, 152, 25, 163, 173, 84, 216, 211, 69, 205, 122, 63,
227, 98, 234, 163, 66, 145, 71, 217, 221, 29, 78, 36, 77, 68, 29, 144,
2, 235, 80, 237, 21, 95, 54, 16, 89, 74, 200, 124, 248, 119, 216, 38,
167, 19, 129, 205, 138, 218, 57, 198, 4, 4, 202, 115, 62, 55, 213, 141,
0,
];

let keys = derive_multiple_phoenix_sk(&SEED, INDEX..INDEX + 2);
assert_eq!(keys[0].to_bytes(), sk_bytes_0,);
assert_eq!(keys[1].to_bytes(), sk_bytes_1,);
}

#[test]
fn test_derive_phoenix_pk() {
// it is important that we always derive the same key from a fixed seed
Expand Down
148 changes: 148 additions & 0 deletions wallet-core/tests/notes.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,148 @@
// 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.

use ff::Field;
use rand::rngs::StdRng;
use rand::SeedableRng;

use execution_core::{
transfer::phoenix::{
Note, PublicKey as PhoenixPublicKey, SecretKey as PhoenixSecretKey,
},
JubJubScalar,
};

use wallet_core::{
keys::derive_multiple_phoenix_sk, map_owned, phoenix_balance, BalanceInfo,
};

#[test]
fn test_map_owned() {
// Assuming this set of notes where the number used as suffix is the
// "owner":
// notes := [A1, B1, C2, D2, E1, F3]

let mut rng = StdRng::seed_from_u64(0xdab);
const SEED_1: [u8; 64] = [1; 64];
const SEED_2: [u8; 64] = [2; 64];

let owner_1_sks = derive_multiple_phoenix_sk(&SEED_1, 0..3);
let owner_1_pks = [
PhoenixPublicKey::from(&owner_1_sks[0]),
PhoenixPublicKey::from(&owner_1_sks[1]),
PhoenixPublicKey::from(&owner_1_sks[2]),
];
let owner_2_sks = derive_multiple_phoenix_sk(&SEED_2, 0..2);
let owner_2_pks = [
PhoenixPublicKey::from(&owner_2_sks[0]),
PhoenixPublicKey::from(&owner_2_sks[1]),
];
let owner_3_pk =
PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng));

let value = 42;
let notes = vec![
gen_note(&mut rng, true, &owner_1_pks[0], value), // owner 1
gen_note(&mut rng, true, &owner_1_pks[1], value), // owner 1
gen_note(&mut rng, true, &owner_2_pks[0], value), // owner 2
gen_note(&mut rng, true, &owner_2_pks[1], value), // owner 2
gen_note(&mut rng, true, &owner_1_pks[2], value), // owner 1
gen_note(&mut rng, true, &owner_3_pk, value), // owner 3
];

// notes with idx 0, 1 and 4 are owned by owner_1
let notes_by_1 = map_owned(&owner_1_sks, &notes);
assert_eq!(notes_by_1.len(), 3);
let note = &notes[0];
let nullifier = note.gen_nullifier(&owner_1_sks[0]);
assert_eq!(&notes_by_1[&nullifier], note);
let note = &notes[1];
let nullifier = note.gen_nullifier(&owner_1_sks[1]);
assert_eq!(&notes_by_1[&nullifier], note);
let note = &notes[4];
let nullifier = note.gen_nullifier(&owner_1_sks[2]);
assert_eq!(&notes_by_1[&nullifier], note);

// notes with idx 2 and 3 are owned by owner_2
let notes_by_2 = map_owned(&owner_2_sks, &notes);
assert_eq!(notes_by_2.len(), 2);
let note = &notes[2];
let nullifier = note.gen_nullifier(&owner_2_sks[0]);
assert_eq!(&notes_by_2[&nullifier], note);
let note = &notes[3];
let nullifier = note.gen_nullifier(&owner_2_sks[1]);
assert_eq!(&notes_by_2[&nullifier], note);
}

#[test]
fn test_balance() {
let mut rng = StdRng::seed_from_u64(0xdab);

let owner_sk = PhoenixSecretKey::random(&mut rng);
let owner_pk = PhoenixPublicKey::from(&owner_sk);

let mut notes = Vec::new();

// create the notes
for value in 0..=10 {
// we want to test with a mix of transparent and obfuscated notes so we
// make every 10th note transparent
let obfuscated_note = if value % 10 == 0 { false } else { true };

notes.push(gen_note(&mut rng, obfuscated_note, &owner_pk, value));

// also push some notes that are not owned
if value % 4 == 0 {
let not_owner_pk =
PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng));
notes.push(gen_note(
&mut rng,
obfuscated_note,
&not_owner_pk,
value,
));
}
}

// the sum of these notes should be 5 * 11 = 55
// and the spendable notes are 7 + 8 + 9 + 10 = 34
let expected_balance = BalanceInfo {
value: 55,
spendable: 34,
};

assert_eq!(
phoenix_balance(&(&owner_sk).into(), notes),
expected_balance
);
}

fn gen_note(
rng: &mut StdRng,
obfuscated_note: bool,
owner_pk: &PhoenixPublicKey,
value: u64,
) -> Note {
let sender_pk = PhoenixPublicKey::from(&PhoenixSecretKey::random(rng));

let value_blinder = JubJubScalar::random(&mut *rng);
let sender_blinder = [
JubJubScalar::random(&mut *rng),
JubJubScalar::random(&mut *rng),
];
if obfuscated_note {
Note::obfuscated(
rng,
&sender_pk,
&owner_pk,
value,
value_blinder,
sender_blinder,
)
} else {
Note::transparent(rng, &sender_pk, &owner_pk, value, sender_blinder)
}
}

0 comments on commit ba594e6

Please sign in to comment.