Mastic aggregator and collector implementation #1107

hannahdaviscrypto · 2024-08-13T18:24:44Z

Implements preparation, aggregation, and unsharding for the Mastic VDAF, and adjusts the VIDPF and SZK modules in support of this development.

Implements aggregator and collector functionality for the Mastic protocol for weighted heavy-hitters and attribute-based metrics.

cjpatton

Comments through prep_init(). Great start!

src/flp/szk.rs

src/vdaf/mastic.rs

src/vidpf.rs

src/flp/szk.rs

src/idpf.rs

src/vdaf/mastic.rs

src/vidpf.rs

cjpatton

Comments up to prepare_init().

src/vidpf.rs

src/vdaf/mastic.rs

src/vidpf.rs

cjpatton · 2024-10-18T20:12:37Z

src/vidpf.rs

    }
 }

 /// Vidpf evaluation state
 ///
 /// Contains the values produced during input evaluation at a given level.
+#[derive(Debug)]


@divergentdave can you check if Debug should be fenced by a feature flag? In the past we discussed using this pattern to prevent secrets from being logged in production environments. I'm not sure what the state of play is now.

We don't currently gate any Debug impls on that feature flag. A number of structs currently derive Debug, like Seed, Prio3InputShare, Poplar1InputShare.

src/vidpf.rs

src/vdaf/mastic.rs

cjpatton · 2024-10-30T22:18:44Z

src/flp/szk.rs

-            part.encode(bytes)?
-        };
-        if let Some(ref part) = self.helper_joint_rand_part_opt {
+        if let Some(ref part) = self.joint_rand_part_opt {
            part.encode(bytes)?


nit: rust convention (AFAIK)

Suggested change

part.encode(bytes)?

part.encode(bytes)?;

cjpatton · 2024-10-30T22:23:10Z

src/flp/szk.rs

+pub type SzkQueryState<const SEED_SIZE: usize> = Option<Seed<SEED_SIZE>>;
+
+/// Verifier type for the SZK proof.
+pub type SzkVerifier<F> = Vec<F>;


It looks like we can now delete this type.

Notably, the two encode/decode implementations are unused, and since they are on a type alias, they would bleed outside of the module to any use of a Vec<T>.

cjpatton · 2024-10-30T22:27:19Z

src/flp/szk.rs

+            .zip(helper_share.flp_verifier)
+        {
+            *x += y;
+        }


General control flow tip: Avoid nesting if statements. E.g., instead of

if cond1 { do_this(); if cond2 { do_that(); } else { handle_error(); } } else { handle_error(); }

do

if !cond1 { handle_error(); } do_this(); if !cond2 { handle_error(); } do_that();

This makes code easier to read and maintain.

cjpatton · 2024-10-30T22:28:47Z

src/flp/szk.rs

+                leader_share.joint_rand_part_opt,
+                helper_share.joint_rand_part_opt,
+            ) {
+                (Some(leader_part), Some(helper_part)) => Ok(Some([leader_part, helper_part])),


Here we could already combine the parts into the seed. The advantage is we reduce the amount of bits on the wire by half. (Plus, IIRC, this is what the current spec says to do.)

cjpatton · 2024-10-30T22:29:25Z

src/flp/szk.rs

+        } else {
+            Err(SzkError::Decide("failed to verify FLP proof".to_string()))
+        }
+    }


Suggested change

}

}

cjpatton · 2024-10-30T23:14:49Z

src/vdaf/mastic.rs

+        }
+    }
+
+    fn aggregate<M: IntoIterator<Item = Self::OutputShare>>(


Convention: Use the concrete type here

Suggested change

fn aggregate<M: IntoIterator<Item = Self::OutputShare>>(

fn aggregate<M: IntoIterator<Item = MasticOutputShare<T::Field>>>(

cjpatton · 2024-10-30T23:18:28Z

src/vdaf/mastic.rs

+        agg_shares: M,
+        _num_measurements: usize,
+    ) -> Result<Self::AggregateResult, VdafError> {
+        let n = agg_param.level_and_prefixes.prefixes().len();


Suggested change

let n = agg_param.level_and_prefixes.prefixes().len();

let num_prefixes = agg_param.level_and_prefixes.prefixes().len();

cjpatton · 2024-10-30T23:22:12Z

src/vdaf/mastic.rs

+        for i in 0..n {
+            let encoded_result = &agg_final.0
+                [i * self.vidpf.weight_parameter..(i + 1) * self.vidpf.weight_parameter];


Use chunks() to get an iterator instead: https://doc.rust-lang.org/std/slice/struct.Chunks.html#example

This is easier to read and allows the compiler to do some more optimization.

cjpatton · 2024-10-30T23:22:15Z

src/vdaf/mastic.rs

+            result.push(
+                self.szk
+                    .typ
+                    .decode_result(&self.szk.typ.truncate(encoded_result.to_vec())?[..], 1)?,


You've hard coded the number of measurements to be 1. In fact, this is supposed to be the number of measurements that have the prefix in common. See https://github.com/jimouris/draft-mouris-cfrg-mastic/blob/482903a99a52cfb61ec06d979ade2fc0512e8c3f/poc/mastic.py#L351.

Rather than address it in this PR, I suggest leaving it as a TODO for #947 and taking care of it later.

It looks like we need to add the counter prefix to the encoded measurement, from sharding through to here. Once that's addressed in a follow-up, I think we'll want convenience constructors for Mastic with particular circuits, to take care of constructing the Szk and Vidpf objects, and calculating the VIDPF weight parameter from szk.typ.output_len() + 1.

I also noticed that the spec currently calls truncate during preparation, before producing an output share, as opposed to during unsharding.

cjpatton · 2024-10-30T23:24:07Z

src/vdaf/mastic.rs

+        let first_input = VidpfInput::from_bytes(&[240u8, 0u8, 1u8, 4u8][..]);
+        let second_input = VidpfInput::from_bytes(&[112u8, 0u8, 1u8, 4u8][..]);
+        let third_input = VidpfInput::from_bytes(&[48u8, 0u8, 1u8, 4u8][..]);
+        let fourth_input = VidpfInput::from_bytes(&[32u8, 0u8, 1u8, 4u8][..]);
+        let fifth_input = VidpfInput::from_bytes(&[0u8, 0u8, 1u8, 4u8][..]);
+        let first_prefix = VidpfInput::from_bools(&[false, false, true]);
+        let second_prefix = VidpfInput::from_bools(&[false]);
+        let third_prefix = VidpfInput::from_bools(&[true]);


Put this in an array and reference the indices of the array below.

Suggested change

let first_input = VidpfInput::from_bytes(&[240u8, 0u8, 1u8, 4u8][..]);

let second_input = VidpfInput::from_bytes(&[112u8, 0u8, 1u8, 4u8][..]);

let third_input = VidpfInput::from_bytes(&[48u8, 0u8, 1u8, 4u8][..]);

let fourth_input = VidpfInput::from_bytes(&[32u8, 0u8, 1u8, 4u8][..]);

let fifth_input = VidpfInput::from_bytes(&[0u8, 0u8, 1u8, 4u8][..]);

let first_prefix = VidpfInput::from_bools(&[false, false, true]);

let second_prefix = VidpfInput::from_bools(&[false]);

let third_prefix = VidpfInput::from_bools(&[true]);

let prefixes = [

VidpfInput::from_bytes(&[240u8, 0u8, 1u8, 4u8][..]),

VidpfInput::from_bytes(&[112u8, 0u8, 1u8, 4u8][..]),

... and so on

];

cjpatton · 2024-10-31T17:01:16Z

src/vdaf/mastic.rs

+        let (public_share, input_shares) =
+            mastic.shard(&(first_input.clone(), false), &nonce).unwrap();
+        run_vdaf_prepare(
+            &mastic,
+            &verify_key,
+            &first_agg_param,
+            &nonce,
+            public_share,
+            input_shares,
+        )
+        .unwrap();
+
+        let (public_share, input_shares) = mastic.shard(&(first_input, true), &nonce).unwrap();
+        run_vdaf_prepare(
+            &mastic,
+            &verify_key,
+            &first_agg_param,
+            &nonce,
+            public_share,
+            input_shares,
+        )
+        .unwrap();


I think we should use run_vdaf() to run the whole thing end-to-end, including aggregation and collection:

libprio-rs/src/vdaf.rs

Line 486 in d2fe428

pub fn run_vdaf<V, M, const SEED_SIZE: usize>(

divergentdave

Partial review, of the SZK module:

divergentdave · 2024-10-22T20:53:02Z

src/vidpf.rs

    }
 }

 /// Vidpf evaluation state
 ///
 /// Contains the values produced during input evaluation at a given level.
+#[derive(Debug)]


We don't currently gate any Debug impls on that feature flag. A number of structs currently derive Debug, like Seed, Prio3InputShare, Poplar1InputShare.

divergentdave · 2024-11-07T16:06:54Z

src/flp/szk.rs

+    /// Returned when a user fails to store the length of the verifier so it
+    /// can be properly read upon receipt.
+    #[error("Part of Szk query state not stored")]
+    InvalidState(String),


This error variant is no longer used for this purpose. It's now used when the leader share and helper share have a mismatched presence/absence of joint randomness parts. Since that error condition is not related to the prepare state, I think we can remove this InvalidState variant.

divergentdave · 2024-11-07T17:09:55Z

src/flp/szk.rs

+    }
+}
+
+/// Joint share type for the SZK proof.


I think this needs some further explanation, since this is a novel term.

Suggested change

/// Joint share type for the SZK proof.

/// Joint share type for the SZK proof.

///

/// This is produced by [`Szk::merge_verifiers`] as the result of combining two query shares.

/// It contains the joint randomness parts from each aggregator, if applicable. It is consumed by [`Szk::decide`].

Or, taking https://github.com/divviup/libprio-rs/pull/1107/files#r1823527498 into account:

Suggested change

/// Joint share type for the SZK proof.

/// Joint share type for the SZK proof.

///

/// This is produced by [`Szk::merge_verifiers`] as the result of combining two query shares.

/// It contains the re-computed joint randomness seed, if applicable. It is consumed by [`Szk::decide`].

divergentdave · 2024-11-07T17:10:37Z

src/flp/szk.rs

-        if !check_flp_proof {
-            return Ok(false);
-        }
+        joint_share: &[Seed<SEED_SIZE>; 2],


I think we could better encapsulate the SZK logic if this took in a &SzkJointShare instead. This would let us remove the match block from Mastic::prepare_next().

divergentdave · 2024-11-07T17:32:38Z

src/flp/szk.rs

+            ) {
+                (Some(leader_part), Some(helper_part)) => Ok(Some([leader_part, helper_part])),
+                (None, None) => Ok(None),
+                _ => Err(SzkError::InvalidState(


As mentioned above, this error case is no longer about the prepare state, since this just depends on query shares.

divergentdave · 2024-11-07T17:50:26Z

src/vdaf/mastic.rs

@@ -5,19 +5,22 @@
 //! [draft-mouris-cfrg-mastic-01]: https://www.ietf.org/archive/id/draft-mouris-cfrg-mastic-01.html

 use crate::{
+    bt::{BinaryTree, Path},


Can we remove #![allow(dead_code)] etc. from the bt module now that this is using it?

Relatedly, I think we ought to make the bt module public now, since BinaryTree appears in the arguments of the public Vidpf::eval_with_cache() method.

divergentdave · 2024-11-07T23:44:30Z

src/vidpf.rs

+        if cache_tree.root.is_none() {
+            cache_tree.root = Some(Box::new(Node::new(VidpfEvalCache {
+                state: VidpfEvalState::init_from_key(id, key),
+                share: W::zero(&self.weight_parameter), // not used
+            })));
+        }
+
+        let mut sub_tree = cache_tree.root.as_mut().expect("root was visited");


We can combine these and remove the .expect() by using Option::get_or_insert_with()

Suggested change

if cache_tree.root.is_none() {

cache_tree.root = Some(Box::new(Node::new(VidpfEvalCache {

state: VidpfEvalState::init_from_key(id, key),

share: W::zero(&self.weight_parameter), // not used

})));

}

let mut sub_tree = cache_tree.root.as_mut().expect("root was visited");

let mut sub_tree = cache_tree.root.get_or_insert_with(|| {

Box::new(Node::new(VidpfEvalCache {

state: VidpfEvalState::init_from_key(id, key),

share: W::zero(&self.weight_parameter), // not used

}))

});

divergentdave · 2024-11-08T00:37:30Z

src/vidpf.rs

+        }
+        let mut cs = Vec::<VidpfProof>::with_capacity(*bits);
+        for _ in 0..*bits {
+            let mut proof = [0u8; 32];


nit: another place where we can use this constant

Suggested change

let mut proof = [0u8; 32];

let mut proof = [0u8; VIDPF_PROOF_SIZE];

divergentdave · 2024-11-08T00:46:30Z

src/vidpf.rs

+            let mut bytes = vec![];
+            public.encode(&mut bytes).unwrap();
+
+            assert_eq!(public.encoded_len().unwrap(), bytes.len());
+
+            let decoded = VidpfPublicShare::<TestWeight>::decode_with_param(
+                &(8, TEST_WEIGHT_LEN),
+                &mut Cursor::new(&bytes),
+            )
+            .unwrap();


nit: Using the get_... helper methods can save some boilerplate, and gets us a check for unused bytes for free.

Suggested change

let mut bytes = vec![];

public.encode(&mut bytes).unwrap();

assert_eq!(public.encoded_len().unwrap(), bytes.len());

let decoded = VidpfPublicShare::<TestWeight>::decode_with_param(

&(8, TEST_WEIGHT_LEN),

&mut Cursor::new(&bytes),

)

.unwrap();

let bytes = public.get_encode().unwrap();

assert_eq!(public.encoded_len().unwrap(), bytes.len());

let decoded = VidpfPublicShare::<TestWeight>::get_decoded_with_param(

&(8, TEST_WEIGHT_LEN),

&bytes,

)

.unwrap();

divergentdave · 2024-11-08T01:19:10Z

src/vdaf/mastic.rs

+/// Mastic prepare state.
+///
+/// State held by an aggregator between rounds of Mastic preparation. Includes intermediate
+/// state for [`Szk``] verification, the output shares currently being validated, and


nit: typo

Suggested change

/// state for [`Szk``] verification, the output shares currently being validated, and

/// state for [`Szk`] verification, the output shares currently being validated, and

divergentdave · 2024-11-08T01:26:13Z

src/vidpf.rs

@@ -226,6 +225,57 @@ impl<W: VidpfValue, const NONCE_SIZE: usize> Vidpf<W, NONCE_SIZE> {
        })
    }

+    /// [`Vidpf::eval_with_cache`] evaluates the entire `input` and produces a share of the


Here and elsewhere, we can drop the name of methods from their documentation, since this will appear just below the full method signature.

Suggested change

/// [`Vidpf::eval_with_cache`] evaluates the entire `input` and produces a share of the

/// Evaluates the entire `input` and produces a share of the

divergentdave · 2024-11-08T01:43:52Z

src/vdaf/mastic.rs

+        let second_prefix = VidpfInput::from_bools(&[false]);
+        let third_prefix = VidpfInput::from_bools(&[true]);
+        let mastic = Mastic::new(algorithm_id, szk, sum_vidpf, 32);
+        let first_agg_param = MasticAggregationParam::new(vec![first_prefix], true).unwrap();


We should add some test cases covering require_weight_check = false.

divergentdave · 2024-11-08T16:19:06Z

src/vdaf/mastic.rs

+
+/// Mastic prepare share.
+///
+/// Broadcast message from an aggregator between rounds of Mastic. Includes the


Here's a suggested rephrasing

Suggested change

/// Broadcast message from an aggregator between rounds of Mastic. Includes the

/// Broadcast message from an aggregator during preparation. Includes the

divergentdave · 2024-11-08T19:40:47Z

src/vdaf/mastic.rs

+            None
+        };
+        let (prep_share, prep_state) =
+            if let Some((szk_query_share, szk_query_state)) = szk_verify_opt {


I think this code could be further simplified by fusing this if-else and the if-else above. szk_verify_opt is constructed by the previous if-else block, and immediately destructured here, with no other uses

divergentdave · 2024-11-08T20:01:26Z

src/vdaf/mastic.rs

+            result.push(
+                self.szk
+                    .typ
+                    .decode_result(&self.szk.typ.truncate(encoded_result.to_vec())?[..], 1)?,


It looks like we need to add the counter prefix to the encoded measurement, from sharding through to here. Once that's addressed in a follow-up, I think we'll want convenience constructors for Mastic with particular circuits, to take care of constructing the Szk and Vidpf objects, and calculating the VIDPF weight parameter from szk.typ.output_len() + 1.

divergentdave · 2024-11-08T20:14:41Z

src/vdaf/mastic.rs

+            result.push(
+                self.szk
+                    .typ
+                    .decode_result(&self.szk.typ.truncate(encoded_result.to_vec())?[..], 1)?,


I also noticed that the spec currently calls truncate during preparation, before producing an output share, as opposed to during unsharding.

hannahdaviscrypto requested a review from a team as a code owner August 13, 2024 18:24

hannahdaviscrypto force-pushed the mastic-agg-dev branch from 5a8d144 to 2c36004 Compare August 13, 2024 19:06

Hannah Davis added 2 commits August 13, 2024 19:06

Updating SZK module for compatibility with Mastic

c67ecb7

Create Mastic module with server implementation

2dfd610

Implements aggregator and collector functionality for the Mastic protocol for weighted heavy-hitters and attribute-based metrics.

hannahdaviscrypto force-pushed the mastic-agg-dev branch from 2c36004 to d7d7fb8 Compare August 13, 2024 19:15

test flag keeps sneaking back in

f952b4e

hannahdaviscrypto force-pushed the mastic-agg-dev branch from d7d7fb8 to f952b4e Compare August 13, 2024 19:18

better encoding for SzkQueryState

96b1c4c

cjpatton self-requested a review August 15, 2024 19:37

cjpatton reviewed Aug 15, 2024

View reviewed changes

first pass at review comments, mainly eliminating round enums

7995ed1

hannahdaviscrypto force-pushed the mastic-agg-dev branch from fa2c390 to 7995ed1 Compare August 27, 2024 19:00

Hannah Davis added 2 commits September 13, 2024 16:36

non-compiling caching edits to vidpf

b9ebfa3

removing fix

130622a

hannahdaviscrypto force-pushed the mastic-agg-dev branch 3 times, most recently from e66a7fb to 69749ff Compare September 16, 2024 20:54

more efficient caching

27d2b2f

hannahdaviscrypto force-pushed the mastic-agg-dev branch from 69749ff to 27d2b2f Compare September 16, 2024 20:58

cjpatton self-requested a review September 17, 2024 17:21

cjpatton reviewed Sep 17, 2024

View reviewed changes

Hannah Davis added 2 commits September 24, 2024 21:51

addressing recent review

5699e50

caching completed

4d47513

hannahdaviscrypto force-pushed the mastic-agg-dev branch from dafe718 to 4d47513 Compare October 3, 2024 03:28

hannahdaviscrypto requested a review from cjpatton October 3, 2024 20:02

cjpatton reviewed Oct 8, 2024

View reviewed changes

remove aggregator ID from VidpfKey

d927640

hannahdaviscrypto requested a review from cjpatton October 17, 2024 23:58

cjpatton reviewed Oct 18, 2024

View reviewed changes

hannahdaviscrypto force-pushed the mastic-agg-dev branch from e0d95bd to 0ea57b9 Compare October 22, 2024 21:37

polishing comments

5a46f32

hannahdaviscrypto force-pushed the mastic-agg-dev branch from 0ea57b9 to 5a46f32 Compare October 22, 2024 21:38

hannahdaviscrypto force-pushed the mastic-agg-dev branch from 6c6e284 to 7ecf3ca Compare October 30, 2024 21:28

refactoring szk

4f487df

hannahdaviscrypto force-pushed the mastic-agg-dev branch from 7ecf3ca to 4f487df Compare October 30, 2024 21:32

fix clippy complaints

e5520fb

cjpatton self-requested a review October 30, 2024 22:08

cjpatton reviewed Oct 30, 2024

View reviewed changes

cjpatton reviewed Oct 31, 2024

View reviewed changes

divergentdave reviewed Nov 7, 2024

View reviewed changes

divergentdave mentioned this pull request Nov 8, 2024

Add a method for is_valid to the Aggregator trait #1133

Closed

divergentdave reviewed Nov 8, 2024

View reviewed changes

	fn aggregate<M: IntoIterator<Item = Self::OutputShare>>(
	fn aggregate<M: IntoIterator<Item = MasticOutputShare<T::Field>>>(

	let n = agg_param.level_and_prefixes.prefixes().len();
	let num_prefixes = agg_param.level_and_prefixes.prefixes().len();

	let mut proof = [0u8; 32];
	let mut proof = [0u8; VIDPF_PROOF_SIZE];

	/// state for [`Szk``] verification, the output shares currently being validated, and
	/// state for [`Szk`] verification, the output shares currently being validated, and

	/// [`Vidpf::eval_with_cache`] evaluates the entire `input` and produces a share of the
	/// Evaluates the entire `input` and produces a share of the

	/// Broadcast message from an aggregator between rounds of Mastic. Includes the
	/// Broadcast message from an aggregator during preparation. Includes the

Mastic aggregator and collector implementation #1107

Are you sure you want to change the base?

Mastic aggregator and collector implementation #1107

Conversation

hannahdaviscrypto commented Aug 13, 2024

cjpatton left a comment

Choose a reason for hiding this comment

cjpatton left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

divergentdave left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

divergentdave Nov 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

divergentdave Nov 7, 2024 •

edited

Loading