From 0d35571215c009cafcc374010188277c6db43501 Mon Sep 17 00:00:00 2001 From: Aaryamann Challani <43716372+rymnc@users.noreply.github.com> Date: Thu, 8 Jun 2023 21:33:09 +0530 Subject: [PATCH] feat(rln): get_leaf ffi and public api (#177) --- rln/src/ffi.rs | 6 ++++++ rln/src/public.rs | 53 +++++++++++++++++++++++++++++++++++++++++++++++ rln/tests/ffi.rs | 45 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+) diff --git a/rln/src/ffi.rs b/rln/src/ffi.rs index f20e205..ab2b9ff 100644 --- a/rln/src/ffi.rs +++ b/rln/src/ffi.rs @@ -222,6 +222,12 @@ pub extern "C" fn set_leaf(ctx: *mut RLN, index: usize, input_buffer: *const Buf call!(ctx, set_leaf, index, input_buffer) } +#[allow(clippy::not_unsafe_ptr_arg_deref)] +#[no_mangle] +pub extern "C" fn get_leaf(ctx: *mut RLN, index: usize, output_buffer: *mut Buffer) -> bool { + call_with_output_arg!(ctx, get_leaf, output_buffer, index) +} + #[allow(clippy::not_unsafe_ptr_arg_deref)] #[no_mangle] pub extern "C" fn set_next_leaf(ctx: *mut RLN, input_buffer: *const Buffer) -> bool { diff --git a/rln/src/public.rs b/rln/src/public.rs index 3f3899f..90d679d 100644 --- a/rln/src/public.rs +++ b/rln/src/public.rs @@ -217,6 +217,31 @@ impl RLN<'_> { Ok(()) } + /// Gets a leaf value at position index in the internal Merkle tree. + /// The leaf value is written to output_data. + /// Input values are: + /// - `index`: the index of the leaf + /// + /// Example: + /// ``` + /// use crate::protocol::*; + /// use std::io::Cursor; + /// + /// let id_index = 10; + /// let mut buffer = Cursor::new(Vec::::new()); + /// rln.get_leaf(id_index, &mut buffer).unwrap(); + /// let id_commitment = deserialize_field_element(&buffer.into_inner()).unwrap(); + pub fn get_leaf(&self, index: usize, mut output_data: W) -> Result<()> { + // We get the leaf at input index + let leaf = self.tree.get(index)?; + + // We serialize the leaf and write it to output + let leaf_byte = fr_to_bytes_le(&leaf); + output_data.write_all(&leaf_byte)?; + + Ok(()) + } + /// Sets multiple leaves starting from position index in the internal Merkle tree. /// /// If n leaves are passed as input, these will be set at positions `index`, `index+1`, ..., `index+n-1` respectively. @@ -1829,4 +1854,32 @@ mod test { // We ensure that an empty value was written to output_buffer, i.e. no secret is recovered assert!(serialized_identity_secret_hash.is_empty()); } + + #[test] + fn test_get_leaf() { + // We generate a random tree + let tree_height = 10; + let mut rng = thread_rng(); + let input_buffer = + Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string()); + let mut rln = RLN::new(tree_height, input_buffer).unwrap(); + + // We generate a random leaf + let leaf = Fr::rand(&mut rng); + + // We generate a random index + let index = rng.gen_range(0..rln.tree.capacity()); + + // We add the leaf to the tree + let mut buffer = Cursor::new(fr_to_bytes_le(&leaf)); + rln.set_leaf(index, &mut buffer).unwrap(); + + // We get the leaf + let mut output_buffer = Cursor::new(Vec::::new()); + rln.get_leaf(index, &mut output_buffer).unwrap(); + + // We ensure that the leaf is the same as the one we added + let (received_leaf, _) = bytes_le_to_fr(output_buffer.into_inner().as_ref()); + assert_eq!(received_leaf, leaf); + } } diff --git a/rln/tests/ffi.rs b/rln/tests/ffi.rs index 3fb28cd..5a192b6 100644 --- a/rln/tests/ffi.rs +++ b/rln/tests/ffi.rs @@ -1156,4 +1156,49 @@ mod test { assert_eq!(received_hash, expected_hash); } + + #[test] + fn test_get_leaf() { + // We create a RLN instance + let tree_height = TEST_TREE_HEIGHT; + let no_of_leaves = 1 << TEST_TREE_HEIGHT; + + let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit(); + let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string(); + let input_buffer = &Buffer::from(input_config.as_bytes()); + let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr()); + assert!(success, "RLN object creation failed"); + let rln_pointer = unsafe { &mut *rln_pointer.assume_init() }; + + // We generate a new identity tuple from an input seed + let seed_bytes: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]; + let input_buffer = &Buffer::from(seed_bytes); + let mut output_buffer = MaybeUninit::::uninit(); + let success = + seeded_extended_key_gen(rln_pointer, input_buffer, output_buffer.as_mut_ptr()); + assert!(success, "seeded key gen call failed"); + + let output_buffer = unsafe { output_buffer.assume_init() }; + let result_data = <&[u8]>::from(&output_buffer).to_vec(); + let (_, _, _, id_commitment) = deserialize_identity_tuple(result_data); + + // We insert the id_commitment into the tree at a random index + let mut rng = thread_rng(); + let index = rng.gen_range(0..no_of_leaves) as usize; + let leaf = fr_to_bytes_le(&id_commitment); + let input_buffer = &Buffer::from(leaf.as_ref()); + let success = set_leaf(rln_pointer, index, input_buffer); + assert!(success, "set leaf call failed"); + + // We get the leaf at the same index + let mut output_buffer = MaybeUninit::::uninit(); + let success = get_leaf(rln_pointer, index, output_buffer.as_mut_ptr()); + assert!(success, "get leaf call failed"); + let output_buffer = unsafe { output_buffer.assume_init() }; + let result_data = <&[u8]>::from(&output_buffer).to_vec(); + let (received_id_commitment, _) = bytes_le_to_fr(&result_data); + + // We check that the received id_commitment is the same as the one we inserted + assert_eq!(received_id_commitment, id_commitment); + } }