[RFC][oprf][shuffle] OPRF Shuffle using a 2-round 4-message shuffle protocol #809

cryo28 · 2023-10-18T19:12:10Z

New query type
Implementation of the protocol (not sharded)

This is not a final code I'd like to land. But more a preview to gather feedback and comments. In a few places I don't really know what I am doing so I resorted to copy-pasta.

If there is not major problems with the code I'll proceed to writing a bunch of unit tests, docs to package the commit for merging.

1. New query type 2. Implementation of the protocol (not sharded)

martinthomson · 2023-10-18T21:46:06Z

src/protocol/context/oprf.rs

+};
+
+#[derive(Clone)]
+pub struct Context<'a> {


What is the purpose of this code?

martinthomson · 2023-10-18T21:50:09Z

src/protocol/oprf/mod.rs

+    Ok(res)
+}
+
+async fn run_h3<C, L, R>(


There is a lot of code duplication between these functions.

martinthomson · 2023-10-18T21:50:57Z

src/helpers/transport/query/oprf_shuffle.rs

+impl Default for QueryConfig {
+    fn default() -> Self {
+        Self {
+            bk_size: 40,


This could probably be smaller.

with peers in run_h{1,2,3} functions

akoshelev

Haven't finished the whole thing, some initial feedback in place

akoshelev · 2023-10-18T23:06:37Z

src/query/runner/oprf_shuffle.rs

+    }
+}
+
+/// Helps to convince the compiler that things are `Send`. Like `seq_join::assert_send`, but for


We should reuse the existing function

akoshelev · 2023-10-18T23:25:14Z

src/protocol/oprf/mod.rs

+    let iter = std::iter::from_fn(move || Some(OPRFShuffleSingleShare::sample(&mut rng)))
+        .take(batch_size as usize);
+
+    // NOTE: I'd like to return an Iterator from here as there is really no need to allocate batch_size of items.


https://rust-random.github.io/rand/rand/trait.Rng.html#method.sample_iter may help here

akoshelev · 2023-10-18T23:36:54Z

src/protocol/oprf/mod.rs

+    type Output = OPRFShuffleSingleShare;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        *self + *rhs // Relies on Copy


you don't need this comment, rustc will complain if Self is not Copy

this is more generic than the implementation above, so you could switch the bodies of add

akoshelev · 2023-10-18T23:54:01Z

src/protocol/oprf/mod.rs

+    GeneratePi12,
+    GeneratePi23,
+    GeneratePi31,
+    GenerateZ12,


for $Z_{ij}$ tables, all helpers need to generate 2 of them using shared randomess from left and right. That should require only 2 steps.

Likely it is the same for permutations

akoshelev · 2023-10-18T23:56:40Z

src/protocol/oprf/mod.rs

+        }
+    }
+
+    pub fn sample<R: Rng>(rng: &mut R) -> Self {


you may want to implement Standard: Distribution<MyStruct> here

akoshelev · 2023-10-19T00:08:15Z

src/protocol/oprf/mod.rs

+) -> Result<Vec<OPRFShuffleSingleShare>, Error> {
+    let (step_send, step_recv, dir) = match ctx.role() {
+        Role::H2 => (
+            OPRFShuffleStep::TransferCHat1,


you should be able to use the same step for send and receive. Depending on the role, one helper will be only sending data, while recipient consuming it at the other end

akoshelev · 2023-10-19T00:29:59Z

src/protocol/oprf/mod.rs

+    Ok(output)
+}
+
+async fn exchange_c_hat<C: Context, I: IntoIterator<Item = OPRFShuffleSingleShare>>(


It is probably easier not to have this function

akoshelev · 2023-10-19T01:08:12Z

src/protocol/oprf/mod.rs

+        ctx,
+        &OPRFShuffleStep::TransferX2,
+        Direction::Right,
+        x_2.clone(),


clone here feels wrong

akoshelev · 2023-10-19T01:12:04Z

src/protocol/oprf/mod.rs

+    permutation
+}
+
+fn permute(


this relies on the IntoInterator specialization to avoid extra allocations, but I don't see why it can't be a one liner apply(permutation, &mut data) on the caller's side

akoshelev · 2023-10-19T01:16:52Z

src/protocol/oprf/mod.rs

+///
+/// ![Apply steps][apply]
+fn apply<T>(permutation: &[u32], values: &mut [T]) {
+    // NOTE: This is copypasta from crate::protocol::sort


should we lift it out of sort and reuse?

akoshelev · 2023-10-19T01:18:23Z

src/protocol/context/mod.rs

@@ -8,6 +9,7 @@ use std::{num::NonZeroUsize, sync::Arc};

 use async_trait::async_trait;
 pub use malicious::{Context as MaliciousContext, Upgraded as UpgradedMaliciousContext};
+pub use oprf::Context as OPRFContext;


…for OPRFShuffleSingleShare

1. Unified GenerateZ_ij and GeneratePi_ij steps 2. Moved assert_stream_send into a new module create:one_off_fns (and updated a few callsites to use it) 3. Replaced try_join! macro with future:try_join fn calls 4. Unified ExchangeCHat1 and ExchangeCHat2 steps into ExchangeCHat 5. Replaced calls to exchange_c_hat function with inline send/receive 6. Dropped OPRFContext. And used Base context instead. I had to make Base::new public 7. Removed unnecessary x_2.clone() un run_h1 8. impl Distrubution<OPRFShuffleSingleShare> for Standard and usage of rng.sample_iter(Standard) instead of manual OPRFShuffle::sample() 9. Got rid of permute fn. Call apply (shuffle) inline

1. Rewrote implt Add for &OPRFShuffleSingleShare in a way that does not spook clippy 2. Hoisted variables for narrow contexts generating random tables into the callers to avoid allocation of them. Instead the addition of random tables are done by combining iterators if possible 3. Extract the pieces of work common for all 3 helprs (generate pis, generate zs) to the calling function

1. Renamed OPRFShuffleSingleShare -> OPRFShare 2. Extract OPRFShare and all its impls into oprf_share module 3. Used apply (permutation) from protocol/sort/apply.rs

danielmasny

Shuffle looks great to me! I have a couple of suggestions and left them as comments. I haven't checked the new query type aspects of the pr.

Thanks for all your work!

danielmasny · 2023-10-20T18:18:18Z

src/protocol/mod.rs

@@ -6,6 +6,7 @@ pub mod context;
 pub mod dp;
 pub mod ipa;
 pub mod modulus_conversion;
+pub mod oprf;


Suggested change

pub mod oprf;

#[cfg(feature = "ipa-prf")]

pub mod oprf;

I used the ipa-prf feature which requires descriptive gate, but it hasn't landed yet but probably will land soon.

danielmasny · 2023-10-20T18:23:13Z

src/protocol/mod.rs

@@ -6,6 +6,7 @@ pub mod context;
 pub mod dp;
 pub mod ipa;
 pub mod modulus_conversion;
+pub mod oprf;


we should probably not call this oprf since it is just a shuffle and a basic protocol that could be used in different context. It is not related to an oprf other than that we want to use it together with a prf/orpf in our new IPA version.

danielmasny · 2023-10-20T18:28:59Z

src/protocol/oprf/mod.rs

+
+use self::oprf_share::{OPRFShare, OprfBK, OprfF, OprfMK};
+use super::{
+    context::Context, ipa::IPAInputRow, sort::apply::apply as apply_permutation, RecordId,


I don't like that we refer here to the sort crate, we should probably refactor apply permutation as its own basic protocol. When we do the sharding, we anyway need to add a new, sharded version of apply which will be separated from sort.

danielmasny · 2023-10-20T18:29:54Z

src/protocol/oprf/mod.rs

+use ipa_macros::Step;
+use rand::{distributions::Standard, seq::SliceRandom, Rng};
+
+use self::oprf_share::{OPRFShare, OprfBK, OprfF, OprfMK};


I would prefer to not use oprf in the names here because it is confusing, rather something with shuffle

danielmasny · 2023-10-20T18:49:42Z

src/protocol/oprf/mod.rs

+    let ctx_b_hat = ctx.narrow(&OPRFShuffleStep::GenerateBHat);
+    let b_hat: Vec<_> =
+        generate_random_table_solo(batch_size, &ctx_b_hat, Direction::Right).collect();
+


is there a reason why you generate a hat, b hat here and not together with z and pi? I think it would be cleaner to generate them within the same function.

danielmasny · 2023-10-20T19:19:35Z

src/protocol/oprf/mod.rs

+
+fn generate_pseudorandom_permutation<R: Rng>(batch_size: u32, rng: &mut R) -> Vec<u32> {
+    let mut permutation = (0..batch_size).collect::<Vec<_>>();
+    permutation.shuffle(rng);


Do we use the same function: .shuffle in our sorting protocol for generating random permutations from a seed? It seems to me that this could be generated on the fly rather than storing this huge decompressed permutation in memory.

danielmasny · 2023-10-20T19:29:47Z

src/protocol/oprf/oprf_share.rs

+    pub breakdown_key: OprfBK,
+    pub trigger_value: OprfF,
+}
+


Could we implement (Weak)SharedValue for OPRFShare? Then we could operate on SharedValue during the shuffle. (WeakSharedValue is part of #795 (comment) and only requires to implement additions).

danielmasny · 2023-10-20T19:32:22Z

src/protocol/oprf/oprf_share.rs

+
+impl Message for OPRFShare {}


We dont need that when implementing SharedValue

danielmasny · 2023-10-20T19:33:11Z

src/protocol/sort/mod.rs

@@ -1,9 +1,9 @@
+pub mod apply;


see above, we should consider remove this from sort and add it e.g. as a basic protocol

danielmasny · 2023-10-20T19:38:28Z

src/protocol/oprf/mod.rs

+    let mut permutation = (0..batch_size).collect::<Vec<_>>();
+    permutation.shuffle(rng);
+    permutation
+}


There is not a single test. We should test whether this is actually correct, by applying it e.g. to a sorted list, and checking the output for inequality with the input list and then sorting the output again and checking the list with the input list for equality. It should be easy to do by using TestWorld

1. Renamed oprf_shuffle module into oprf::shuffle 2. Renamed OPRFShare to ShuffleShare 3. Renamed OPRFShuffle* type aliases (for fields) into ShufleShare*

…individual basic protocol Moved protocols::sort::apply to protocols::basic:apply_permutation

…module This is as prep for making oprf/shuffle protocol generic

…ying some trait bounds the shuffle, run_hN functions do not depend on a particular format of input rows. Instead they can receiv IntoIterators consuming anything satisfying a few bounds I have also moved the logic of dealing with specific inputs (deserialize, serialize, split into individual shares, combine back) into a submodule of queiry::runner::oprf, as I expect all this code to be thrown away/reworked, when the actual formats of inputs/outputs are more clear. But for now this change should let me focus on writing unit tests for the protocol, that should remain relevant even if the data input/output formats change

…erator of ReplicatedSecretShares Instead of accepting a tuple of 2 iterators returning SharedValues accept a single iterator of AdditiveShares: ReplicatedSecretShares This should make interaction with TestWorld easier

This makes the code slighlty simpler

Instead, shuffle in-place based on shared randomness

If a vec produced by permutation is used only to compute some other "table" by adding some other table to it, we can do that in-place

cryo28 · 2023-10-25T17:16:05Z

Closing in favor of #816

[oprf][shuffle] OPRF Shuffle using a 2-round 4-message shuffle protocol

dae7449

1. New query type 2. Implementation of the protocol (not sharded)

martinthomson reviewed Oct 18, 2023

View reviewed changes

cryo28 force-pushed the oprf-shuffle branch 4 times, most recently from d24e086 to 572bb25 Compare October 18, 2023 23:39

Removed duplication of generating permutations and "Z" random tables

cd6b5d2

with peers in run_h{1,2,3} functions

cryo28 force-pushed the oprf-shuffle branch from 572bb25 to cd6b5d2 Compare October 18, 2023 23:46

akoshelev reviewed Oct 19, 2023

View reviewed changes

Artem Ignatyev added 4 commits October 19, 2023 12:16

swapped bodies for impl Add for &OPRFShuffleSingleShare and impl Add …

f9119bf

…for OPRFShuffleSingleShare

[oprf][shuffle] Renames

dd44876

1. Renamed OPRFShuffleSingleShare -> OPRFShare 2. Extract OPRFShare and all its impls into oprf_share module 3. Used apply (permutation) from protocol/sort/apply.rs

danielmasny reviewed Oct 20, 2023

View reviewed changes

Artem Ignatyev added 8 commits October 20, 2023 14:51

Merge remote-tracking branch 'origin/main' into oprf-shuffle

a53692d

[oprf][shuffle] A bunch of renames

7cba0bd

1. Renamed oprf_shuffle module into oprf::shuffle 2. Renamed OPRFShare to ShuffleShare 3. Renamed OPRFShuffle* type aliases (for fields) into ShufleShare*

[oprf][shuffle] Made apply permutation module from sort protocoal an …

cd59872

…individual basic protocol Moved protocols::sort::apply to protocols::basic:apply_permutation

[oprf][shuffle] Move share splitting and combining logic into runner …

9043ec3

…module This is as prep for making oprf/shuffle protocol generic

[oprf][shuffle] Tests

2ea0ca1

[oprf][shuffle] Represent shuffle response as Vec<AdditiveShare<_>>

1d9138d

This makes the code slighlty simpler

cryo28 force-pushed the oprf-shuffle branch from 5f65a3b to 1d9138d Compare October 23, 2023 21:49

Artem Ignatyev added 3 commits October 25, 2023 09:41

Merge remote-tracking branch 'origin/main' into oprf-shuffle

cf865b0

[oprf][shuffle] Avoid allocation of permutations

7791ae0

Instead, shuffle in-place based on shared randomness

[oprf][shuffle] Removed a few unnecessary allocations

b3b09ed

If a vec produced by permutation is used only to compute some other "table" by adding some other table to it, we can do that in-place

cryo28 closed this Oct 25, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RFC][oprf][shuffle] OPRF Shuffle using a 2-round 4-message shuffle protocol #809

[RFC][oprf][shuffle] OPRF Shuffle using a 2-round 4-message shuffle protocol #809

cryo28 commented Oct 18, 2023 •

edited

Loading

martinthomson Oct 18, 2023

martinthomson Oct 18, 2023

martinthomson Oct 18, 2023

akoshelev left a comment

akoshelev Oct 18, 2023

akoshelev Oct 18, 2023

akoshelev Oct 18, 2023

akoshelev Oct 18, 2023

akoshelev Oct 18, 2023

akoshelev Oct 18, 2023

akoshelev Oct 19, 2023

akoshelev Oct 19, 2023

akoshelev Oct 19, 2023

akoshelev Oct 19, 2023

akoshelev Oct 19, 2023

akoshelev Oct 19, 2023

danielmasny left a comment

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

danielmasny Oct 20, 2023

cryo28 commented Oct 25, 2023 •

edited

Loading

[RFC][oprf][shuffle] OPRF Shuffle using a 2-round 4-message shuffle protocol #809

[RFC][oprf][shuffle] OPRF Shuffle using a 2-round 4-message shuffle protocol #809

Conversation

cryo28 commented Oct 18, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

akoshelev 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

danielmasny 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

cryo28 commented Oct 25, 2023 • edited Loading

cryo28 commented Oct 18, 2023 •

edited

Loading

cryo28 commented Oct 25, 2023 •

edited

Loading