diff --git a/Cargo.lock b/Cargo.lock index 5ccc0bfa..1dca0c40 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1510,7 +1510,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1796,9 +1796,9 @@ checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" [[package]] name = "h2" -version = "0.4.13" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" dependencies = [ "atomic-waker", "bytes", @@ -2255,7 +2255,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3344,7 +3344,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4030,7 +4030,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4686,7 +4686,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] diff --git a/ampc-actor-utils/src/network/mpc/handle/config.rs b/ampc-actor-utils/src/network/mpc/handle/config.rs index 16f5b18c..319a4e59 100644 --- a/ampc-actor-utils/src/network/mpc/handle/config.rs +++ b/ampc-actor-utils/src/network/mpc/handle/config.rs @@ -1,4 +1,4 @@ -use std::{cmp, time::Duration}; +use std::time::Duration; #[derive(Default, Clone, Debug)] pub struct MpcConfig { @@ -11,24 +11,24 @@ pub struct MpcConfig { impl MpcConfig { pub fn new(timeout_duration: Duration, num_connections: usize, num_sessions: usize) -> Self { - // don't allow fewer requests than connections... - let connection_parallelism = cmp::min(num_connections, num_sessions); + assert!(num_connections > 0, "MPC networking requires a connection"); + assert!(num_sessions > 0, "MPC networking requires a session"); Self { timeout_duration, num_sessions: num_sessions as u32, - num_connections: connection_parallelism as u32, + // A logical session is striped over every connection. Do not clamp + // this to the session count: batch-size-one requests need multiple + // physical flows to exceed a cloud provider's per-flow limit. + num_connections: num_connections as u32, } } pub fn get_sessions_for_connection(&self, idx: u32) -> u32 { - let num_sessions = self.num_sessions; - let num_connections = self.num_connections; - num_sessions / num_connections - + if idx < (num_sessions % num_connections) { - 1 - } else { - 0 - } + if idx < self.num_connections { + self.num_sessions + } else { + 0 + } } } diff --git a/ampc-actor-utils/src/network/mpc/handle/control_channel.rs b/ampc-actor-utils/src/network/mpc/handle/control_channel.rs index 96c42afc..c3263f8f 100644 --- a/ampc-actor-utils/src/network/mpc/handle/control_channel.rs +++ b/ampc-actor-utils/src/network/mpc/handle/control_channel.rs @@ -150,19 +150,19 @@ impl ControlChannel for TcpControlChannel { } async fn sync(&mut self) -> Result<()> { - let token = NetworkValue::Bytes(SYNC_TOKEN_BYTES.to_vec()); + let token = NetworkValue::Bytes(SYNC_TOKEN_BYTES.to_vec().into()); self.send_next(token.clone()).await?; self.send_prev(token).await?; let next_token = self.recv_next().await?; match next_token { - NetworkValue::Bytes(ref bytes) if bytes == SYNC_TOKEN_BYTES => {} + NetworkValue::Bytes(ref bytes) if &bytes[..] == SYNC_TOKEN_BYTES => {} _ => bail!("invalid sync token received from next party"), } let prev_token = self.recv_prev().await?; match prev_token { - NetworkValue::Bytes(ref bytes) if bytes == SYNC_TOKEN_BYTES => {} + NetworkValue::Bytes(ref bytes) if &bytes[..] == SYNC_TOKEN_BYTES => {} _ => bail!("invalid sync token received from prev party"), } diff --git a/ampc-actor-utils/src/network/mpc/handle/session/mod.rs b/ampc-actor-utils/src/network/mpc/handle/session/mod.rs index 924eac20..b5ca0bb0 100644 --- a/ampc-actor-utils/src/network/mpc/handle/session/mod.rs +++ b/ampc-actor-utils/src/network/mpc/handle/session/mod.rs @@ -14,22 +14,151 @@ use crate::{ }, }; use async_trait::async_trait; -use eyre::{eyre, Result}; -use std::collections::HashMap; +use bytes::Bytes; +use eyre::{bail, ensure, eyre, Result}; +use std::collections::{BTreeMap, HashMap}; use tokio::{ sync::mpsc::{self}, time::timeout, }; +const STRIPE_MAGIC: &[u8; 8] = b"AMPCSTRP"; +const STRIPE_HEADER_BYTES: usize = 32; +const MIN_BYTES_PER_STRIPE: usize = 256 * 1024; + +#[derive(Debug)] +struct Fragment { + sequence: u64, + total_len: usize, + index: usize, + count: usize, + payload: Bytes, +} + +impl Fragment { + fn encode( + sequence: u64, + total_len: usize, + index: usize, + count: usize, + payload: &[u8], + ) -> Result { + let total_len = u64::try_from(total_len)?; + let index = u32::try_from(index)?; + let count = u32::try_from(count)?; + let mut bytes = Vec::with_capacity(STRIPE_HEADER_BYTES + payload.len()); + bytes.extend_from_slice(STRIPE_MAGIC); + bytes.extend_from_slice(&sequence.to_le_bytes()); + bytes.extend_from_slice(&total_len.to_le_bytes()); + bytes.extend_from_slice(&index.to_le_bytes()); + bytes.extend_from_slice(&count.to_le_bytes()); + bytes.extend_from_slice(payload); + Ok(NetworkValue::Bytes(bytes.into())) + } + + fn decode(value: NetworkValue) -> Result { + let NetworkValue::Bytes(bytes) = value else { + bail!("striped MPC session received an unframed value"); + }; + ensure!( + bytes.len() >= STRIPE_HEADER_BYTES, + "striped MPC fragment is shorter than its header" + ); + ensure!( + &bytes[..STRIPE_MAGIC.len()] == STRIPE_MAGIC, + "invalid striped MPC fragment magic" + ); + + let sequence = u64::from_le_bytes(bytes[8..16].try_into()?); + let total_len = usize::try_from(u64::from_le_bytes(bytes[16..24].try_into()?))?; + let index = u32::from_le_bytes(bytes[24..28].try_into()?) as usize; + let count = u32::from_le_bytes(bytes[28..32].try_into()?) as usize; + ensure!(count > 0, "striped MPC fragment has zero chunks"); + ensure!(index < count, "striped MPC fragment index is out of range"); + ensure!(total_len > 0, "striped MPC fragment has an empty message"); + + Ok(Self { + sequence, + total_len, + index, + count, + payload: bytes.slice(STRIPE_HEADER_BYTES..), + }) + } +} + +#[derive(Debug)] +struct Assembly { + total_len: usize, + chunks: Vec>, + received: usize, + received_len: usize, +} + +impl Assembly { + fn new(fragment: &Fragment) -> Self { + Self { + total_len: fragment.total_len, + chunks: vec![None; fragment.count], + received: 0, + received_len: 0, + } + } + + fn insert(&mut self, fragment: Fragment) -> Result<()> { + ensure!( + self.total_len == fragment.total_len && self.chunks.len() == fragment.count, + "inconsistent striped MPC fragment metadata" + ); + ensure!( + self.chunks[fragment.index].is_none(), + "duplicate striped MPC fragment" + ); + self.received += 1; + self.received_len += fragment.payload.len(); + ensure!( + self.received_len <= self.total_len, + "striped MPC fragments exceed declared message length" + ); + self.chunks[fragment.index] = Some(fragment.payload); + Ok(()) + } + + fn is_complete(&self) -> bool { + self.received == self.chunks.len() + } + + fn assemble(self) -> Result { + ensure!(self.is_complete(), "striped MPC message is incomplete"); + ensure!( + self.received_len == self.total_len, + "striped MPC message length mismatch" + ); + let mut bytes = Vec::with_capacity(self.total_len); + for chunk in self.chunks { + let chunk = chunk.ok_or_else(|| eyre!("missing striped MPC fragment"))?; + bytes.extend_from_slice(&chunk); + } + NetworkValue::deserialize(&bytes) + } +} + +#[derive(Debug)] +struct PeerIo { + identity: Identity, + tx: Vec, + rx: InStream, + next_send_sequence: u64, + next_receive_sequence: u64, + pending: BTreeMap, +} + #[derive(Debug)] pub struct TcpSession { session_id: SessionId, - // TcpSession is typically used with 2 peers, with a short Identity String. - // In this case, it is better to avoid a HashMap. - identities: Vec, - // channels to the peers in identities - tx: Vec, - rx: Vec, + // TcpSession is typically used with two peers. A logical peer stream owns + // all physical connections so one large request can use aggregate bandwidth. + peers: Vec, config: MpcConfig, } @@ -37,17 +166,28 @@ impl TcpSession { pub fn new( session_id: SessionId, identities: Vec, - tx: Vec, + tx: Vec>, rx: Vec, config: MpcConfig, ) -> Self { assert_eq!(identities.len(), tx.len()); assert_eq!(identities.len(), rx.len()); + let peers = identities + .into_iter() + .zip(tx) + .zip(rx) + .map(|((identity, tx), rx)| PeerIo { + identity, + tx, + rx, + next_send_sequence: 0, + next_receive_sequence: 0, + pending: BTreeMap::new(), + }) + .collect(); Self { session_id, - identities, - tx, - rx, + peers, config, } } @@ -56,18 +196,28 @@ impl TcpSession { self.session_id } - fn get_tx(&self, id: &Identity) -> Option<&OutStream> { - self.identities - .iter() - .position(|x| x == id) - .and_then(|idx| self.tx.get(idx)) + fn get_peer_mut(&mut self, id: &Identity) -> Option<&mut PeerIo> { + self.peers.iter_mut().find(|peer| &peer.identity == id) } - fn get_rx(&mut self, id: &Identity) -> Option<&mut InStream> { - self.identities - .iter() - .position(|x| x == id) - .and_then(|idx| self.rx.get_mut(idx)) + fn take_completed(peer: &mut PeerIo) -> Result> { + let sequence = peer.next_receive_sequence; + let is_complete = peer + .pending + .get(&sequence) + .is_some_and(Assembly::is_complete); + if !is_complete { + return Ok(None); + } + let assembly = peer + .pending + .remove(&sequence) + .ok_or_else(|| eyre!("completed striped MPC message disappeared"))?; + peer.next_receive_sequence = peer + .next_receive_sequence + .checked_add(1) + .ok_or_else(|| eyre!("striped MPC receive sequence exhausted"))?; + assembly.assemble().map(Some) } } @@ -80,29 +230,100 @@ impl Drop for TcpSession { #[async_trait] impl Networking for TcpSession { async fn send(&mut self, value: NetworkValue, receiver: &Identity) -> Result<()> { - let outgoing_stream = self.get_tx(receiver).ok_or(eyre!( + let session_id = self.session_id; + let peer = self.get_peer_mut(receiver).ok_or(eyre!( "Outgoing stream for {receiver:?} in session {:?} not found", - self.session_id + session_id ))?; - outgoing_stream - .send((self.session_id, value)) - .map_err(|e| eyre!(e.to_string()))?; + ensure!( + !peer.tx.is_empty(), + "striped MPC session has no connections" + ); + + if peer.tx.len() == 1 { + peer.tx[0] + .send((session_id, value)) + .map_err(|e| eyre!(e.to_string()))?; + return Ok(()); + } + + let serialized = value.to_network(); + let stripe_count = peer + .tx + .len() + .min(serialized.len().div_ceil(MIN_BYTES_PER_STRIPE).max(1)); + let sequence = peer.next_send_sequence; + peer.next_send_sequence = peer + .next_send_sequence + .checked_add(1) + .ok_or_else(|| eyre!("striped MPC send sequence exhausted"))?; + let first_connection = usize::try_from(sequence % peer.tx.len() as u64)?; + let stripe_base_len = serialized.len() / stripe_count; + let extra_bytes = serialized.len() % stripe_count; + + for index in 0..stripe_count { + let start = index * stripe_base_len + index.min(extra_bytes); + let end = start + stripe_base_len + usize::from(index < extra_bytes); + let payload = &serialized[start..end]; + let connection = (first_connection + index) % peer.tx.len(); + let fragment = + Fragment::encode(sequence, serialized.len(), index, stripe_count, payload)?; + peer.tx[connection] + .send((session_id, fragment)) + .map_err(|e| eyre!(e.to_string()))?; + } Ok(()) } async fn receive(&mut self, sender: &Identity) -> Result { let session_id = self.session_id; let timeout_duration = self.config.timeout_duration; - let incoming_stream = self.get_rx(sender).ok_or(eyre!( + let max_fragments = self.config.num_connections as usize; + let peer = self.get_peer_mut(sender).ok_or(eyre!( "Incoming stream for {sender:?} in session {:?} not found", session_id ))?; - match timeout(timeout_duration, incoming_stream.recv()).await { - Ok(res) => res.ok_or(eyre!("No message received")), + if peer.tx.len() == 1 { + return match timeout(timeout_duration, peer.rx.recv()).await { + Ok(res) => res.ok_or_else(|| eyre!("No message received")), + Err(_) => Err(eyre!( + "Timeout while waiting for message from {sender:?} in {:?}", + session_id + )), + }; + } + + match timeout(timeout_duration, async { + loop { + if let Some(value) = Self::take_completed(peer)? { + return Ok(value); + } + let value = peer + .rx + .recv() + .await + .ok_or_else(|| eyre!("No message received"))?; + let fragment = Fragment::decode(value)?; + ensure!( + fragment.count <= max_fragments, + "striped MPC message exceeds configured connection count" + ); + ensure!( + fragment.sequence >= peer.next_receive_sequence, + "received a stale striped MPC message" + ); + peer.pending + .entry(fragment.sequence) + .or_insert_with(|| Assembly::new(&fragment)) + .insert(fragment)?; + } + }) + .await + { + Ok(res) => res, Err(_) => Err(eyre!( - "Timeout while waiting for message from {sender:?} in \ - {:?}", - self.session_id + "Timeout while waiting for striped message from {sender:?} in {:?}", + session_id )), } } @@ -206,16 +427,22 @@ async fn make_sessions_inner( { let mut tx = Vec::with_capacity(peer_ids.len()); let mut rx = Vec::with_capacity(peer_ids.len()); - let connection_id = ConnectionId::from(idx as u32 % num_connections); for peer_id in &peer_ids { - let outbound_tx = sc - .outbound_tx - .get(peer_id) - .unwrap() - .get(&connection_id) - .cloned() - .unwrap(); + // Rotate the starting connection per session for small messages; + // large messages use this entire vector concurrently. + let first_connection = idx as u32 % num_connections; + let outbound_tx = (0..num_connections) + .map(|offset| ConnectionId::from((first_connection + offset) % num_connections)) + .map(|connection_id| { + sc.outbound_tx + .get(peer_id) + .unwrap() + .get(&connection_id) + .cloned() + .unwrap() + }) + .collect(); tx.push(outbound_tx); let inbound_rx = sc .inbound_rx @@ -233,3 +460,118 @@ async fn make_sessions_inner( sessions } + +#[cfg(test)] +mod tests { + use super::*; + use ampc_secret_sharing::shares::ring_impl::RingElement; + use std::time::Duration; + use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; + + fn test_session( + connections: usize, + ) -> ( + TcpSession, + Identity, + Vec>, + UnboundedSender, + ) { + let identity = Identity::from("peer"); + let mut tx = Vec::with_capacity(connections); + let mut outbound_rx = Vec::with_capacity(connections); + for _ in 0..connections { + let (connection_tx, connection_rx) = mpsc::unbounded_channel(); + tx.push(connection_tx); + outbound_rx.push(connection_rx); + } + let (inbound_tx, inbound_rx) = mpsc::unbounded_channel(); + let config = MpcConfig::new(Duration::from_secs(1), connections, 1); + let session = TcpSession::new( + SessionId::from(7), + vec![identity.clone()], + vec![tx], + vec![inbound_rx], + config, + ); + (session, identity, outbound_rx, inbound_tx) + } + + fn drain_fragments( + outbound: &mut [UnboundedReceiver], + ) -> Vec<(u64, NetworkValue)> { + let mut fragments = Vec::new(); + for connection in outbound { + while let Ok((session_id, value)) = connection.try_recv() { + assert_eq!(session_id, SessionId::from(7)); + let sequence = Fragment::decode(value.clone()).unwrap().sequence; + fragments.push((sequence, value)); + } + } + fragments + } + + #[tokio::test] + async fn one_session_uses_all_connections_for_a_large_value() -> Result<()> { + let (mut session, peer, mut outbound, _inbound) = test_session(4); + let value = NetworkValue::VecRing16(vec![RingElement(42); 1_000_000]); + session.send(value, &peer).await?; + + let fragments = drain_fragments(&mut outbound); + assert_eq!(fragments.len(), 4); + assert!(outbound.iter().all(UnboundedReceiver::is_empty)); + Ok(()) + } + + #[tokio::test] + async fn one_connection_uses_direct_unframed_transport() -> Result<()> { + let (mut session, peer, mut outbound, inbound) = test_session(1); + let value = NetworkValue::VecRing16(vec![RingElement(42); 1_000_000]); + + session.send(value.clone(), &peer).await?; + let (session_id, sent) = outbound[0].recv().await.unwrap(); + assert_eq!(session_id, SessionId::from(7)); + assert_eq!(sent, value); + assert!(Fragment::decode(sent).is_err()); + + inbound.send(value.clone())?; + assert_eq!(session.receive(&peer).await?, value); + Ok(()) + } + + #[tokio::test] + async fn out_of_order_stripes_preserve_logical_message_order() -> Result<()> { + let (mut session, peer, mut outbound, inbound) = test_session(4); + let first = NetworkValue::Bytes(vec![1; 1024 * 1024].into()); + let second = NetworkValue::Bytes(vec![2; 1024 * 1024].into()); + session.send(first.clone(), &peer).await?; + session.send(second.clone(), &peer).await?; + + let fragments = drain_fragments(&mut outbound); + assert_eq!(fragments.len(), 8); + for (_, value) in fragments + .iter() + .filter(|(sequence, _)| *sequence == 1) + .rev() + { + inbound.send(value.clone())?; + } + for (_, value) in fragments + .iter() + .filter(|(sequence, _)| *sequence == 0) + .rev() + { + inbound.send(value.clone())?; + } + + assert_eq!(session.receive(&peer).await?, first); + assert_eq!(session.receive(&peer).await?, second); + Ok(()) + } + + #[test] + fn connection_count_is_not_clamped_to_session_count() { + let config = MpcConfig::new(Duration::from_secs(1), 16, 1); + assert_eq!(config.num_connections, 16); + assert_eq!(config.get_sessions_for_connection(15), 1); + } +} diff --git a/ampc-actor-utils/src/network/mpc/value.rs b/ampc-actor-utils/src/network/mpc/value.rs index d690b6f5..224c76e8 100644 --- a/ampc-actor-utils/src/network/mpc/value.rs +++ b/ampc-actor-utils/src/network/mpc/value.rs @@ -1,7 +1,7 @@ use ampc_secret_sharing::shares::{ self, bit::Bit, ring48::Ring48, ring_impl::RingElement, IntRing2k, }; -use bytes::BytesMut; +use bytes::{Bytes, BytesMut}; use eyre::{bail, eyre, Result}; use num_enum::{IntoPrimitive, TryFromPrimitive}; use std::mem::size_of; @@ -74,7 +74,7 @@ pub enum NetworkValue { NetworkVec(Vec), // used to verify that the PRFs aren't out of sync PrfCheck(RingElement), - Bytes(Vec), + Bytes(Bytes), /// Packed bit vector: (packed_bytes, bit_count) /// Each byte contains 8 bits in LSB-first order VecRingBit(Vec, usize), @@ -429,7 +429,9 @@ impl NetworkValue { 5 + len ); } - Ok(NetworkValue::Bytes(serialized[5..5 + len].to_vec())) + Ok(NetworkValue::Bytes(Bytes::copy_from_slice( + &serialized[5..5 + len], + ))) } DescriptorByte::VecRingBit => { if serialized.len() < 5 { diff --git a/ampc-actor-utils/src/protocol/binary.rs b/ampc-actor-utils/src/protocol/binary.rs index 3ee8b326..ea2bb053 100644 --- a/ampc-actor-utils/src/protocol/binary.rs +++ b/ampc-actor-utils/src/protocol/binary.rs @@ -29,6 +29,17 @@ struct VecBinShare { inner: VecShare, } +type PackedBitPlanes = Vec>; +type PackedBitPlanePair = (PackedBitPlanes, PackedBitPlanes); +type PackedAdditiveComponents = (PackedBitPlanes, PackedBitPlanes, PackedBitPlanes); + +struct ThreeWayAdderParts { + s: PackedBitPlanes, + x1x3: PackedBitPlanes, + x2x3: PackedBitPlanes, + x3: PackedBitPlanes, +} + impl VecBinShare { fn from_ab(a: Vec>, b: Vec>) -> Self { Self { @@ -224,6 +235,42 @@ where Ok(complete_shares) } +/// Evaluate several independent packed AND groups in one communication +/// round, then restore their original group boundaries. +async fn and_many_grouped( + session: &mut Session, + groups: Vec<(SliceShare<'_, T>, SliceShare<'_, T>)>, +) -> Result>, Error> +where + T: NetworkInt + RingRandFillable, + Standard: Distribution, +{ + let mut sizes = Vec::with_capacity(groups.len()); + let mut total = 0; + for (left, right) in &groups { + if left.len() != right.len() { + bail!("InvalidSize in and_many_grouped"); + } + sizes.push(left.len()); + total += left.len(); + } + + let left = groups.iter().flat_map(|(left, _)| left.iter().copied()); + let right = groups.iter().flat_map(|(_, right)| right.iter().copied()); + let local = and_many_iter_send(session, left, right, total).await?; + let remote = and_many_receive(session).await?; + if local.len() != total || remote.len() != total { + bail!("InvalidSize returned by grouped AND"); + } + + let mut local = local.into_iter(); + let mut remote = remote.into_iter(); + Ok(sizes + .into_iter() + .map(|size| VecShare::from_iter_ab(local.by_ref().take(size), remote.by_ref().take(size))) + .collect()) +} + /// Reduce the given vector of bit-vector shares by computing their element-wise AND. /// /// Each vector in `v` is expected to have `len` bits. @@ -390,6 +437,245 @@ where Ok((res1, res2)) } +/// Evaluate the two independent adders used by the fixed anonymous threshold +/// in lockstep. This preserves the exact gates and traffic of the individual +/// 16-bit mask-carry and 18-bit expression adders, but combines equal-depth +/// AND layers into 17 communication rounds instead of running 35 rounds +/// sequentially. +async fn binary_add_anon_threshold_fused( + session: &mut Session, + mask_x1: Vec>, + mask_x2: Vec>, + mask_x3: Vec>, + expression_x1: Vec>, + expression_x2: Vec>, + expression_x3: Vec>, +) -> Result, Error> { + const MASK_BITS: usize = 16; + const EXPRESSION_BITS: usize = 18; + if mask_x1.len() != MASK_BITS + || mask_x2.len() != MASK_BITS + || mask_x3.len() != MASK_BITS + || expression_x1.len() != EXPRESSION_BITS + || expression_x2.len() != EXPRESSION_BITS + || expression_x3.len() != EXPRESSION_BITS + { + bail!("invalid bit width for fused anonymous-threshold adder"); + } + + // Reduce each three-input addition to 2*c+s. Their initial AND layers are + // independent, so send all 16+17 packed gates together. + let mut mask_x2x3 = mask_x2; + transposed_pack_xor_assign(&mut mask_x2x3, &mask_x3); + let mask_s = transposed_pack_xor(&mask_x1, &mask_x2x3); + let mut mask_x1x3 = mask_x1; + transposed_pack_xor_assign(&mut mask_x1x3, &mask_x3); + + let mut expression_x2x3 = expression_x2; + transposed_pack_xor_assign(&mut expression_x2x3, &expression_x3); + let expression_s = transposed_pack_xor(&expression_x1, &expression_x2x3); + let mut expression_x1x3 = expression_x1; + transposed_pack_xor_assign(&mut expression_x1x3, &expression_x3); + + let mask = ThreeWayAdderParts { + s: mask_s, + x1x3: mask_x1x3, + x2x3: mask_x2x3, + x3: mask_x3, + }; + let expression = ThreeWayAdderParts { + s: expression_s, + x1x3: expression_x1x3, + x2x3: expression_x2x3, + x3: expression_x3, + }; + binary_add_anon_threshold_from_parts(session, mask, expression).await +} + +/// Finish the anonymous-threshold circuit from the XOR sum and majority-gate +/// inputs of its two three-way additions. +async fn binary_add_anon_threshold_from_parts( + session: &mut Session, + mut mask: ThreeWayAdderParts, + mut expression: ThreeWayAdderParts, +) -> Result, Error> { + const MASK_BITS: usize = 16; + const EXPRESSION_BITS: usize = 18; + if mask.s.len() != MASK_BITS + || mask.x1x3.len() != MASK_BITS + || mask.x2x3.len() != MASK_BITS + || mask.x3.len() != MASK_BITS + || expression.s.len() != EXPRESSION_BITS + || expression.x1x3.len() != EXPRESSION_BITS + || expression.x2x3.len() != EXPRESSION_BITS + || expression.x3.len() != EXPRESSION_BITS + { + bail!("invalid bit width for fused anonymous-threshold adder"); + } + + expression.x1x3.pop(); + expression.x2x3.pop(); + let mut expression_x3_without_msb = expression.x3; + expression_x3_without_msb.pop(); + + let initial_groups = mask + .x1x3 + .iter() + .zip(&mask.x2x3) + .chain(expression.x1x3.iter().zip(&expression.x2x3)) + .map(|(left, right)| (left.as_slice(), right.as_slice())) + .collect(); + let mut carries = and_many_grouped(session, initial_groups).await?; + let mut expression_c = carries.split_off(MASK_BITS); + let mut mask_c = carries; + transposed_pack_xor_assign(&mut mask_c, &mask.x3); + transposed_pack_xor_assign(&mut expression_c, &expression_x3_without_msb); + let mask_c_msb = mask_c.pop().ok_or_else(|| eyre!("missing mask carry"))?; + + // First ripple gate for each adder. + let mut first = and_many_grouped( + session, + vec![ + (mask.s[1].as_slice(), mask_c[0].as_slice()), + (expression.s[1].as_slice(), expression_c[0].as_slice()), + ], + ) + .await?; + let mut mask_carry = first.remove(0); + let mut expression_carry = first.remove(0); + + // Fourteen carry layers are shared by both adders. + for offset in 0..14 { + let mask_s_bit = &mut mask.s[offset + 2]; + let mask_c_bit = &mut mask_c[offset + 1]; + *mask_s_bit ^= mask_carry.as_slice(); + *mask_c_bit ^= mask_carry.as_slice(); + + let expression_s_bit = &mut expression.s[offset + 2]; + let expression_c_bit = &mut expression_c[offset + 1]; + *expression_s_bit ^= expression_carry.as_slice(); + *expression_c_bit ^= expression_carry.as_slice(); + + let mut next = and_many_grouped( + session, + vec![ + (mask_s_bit.as_slice(), mask_c_bit.as_slice()), + (expression_s_bit.as_slice(), expression_c_bit.as_slice()), + ], + ) + .await?; + mask_carry ^= next.remove(0); + expression_carry ^= next.remove(0); + } + + // The expression has one final ripple layer. In the same round, compute + // the mask's top overflow carry and the final correction AND, whose inputs + // are already available locally at this depth. + let mut mask_carry_16 = mask_c_msb.clone(); + mask_carry_16 ^= mask_carry.as_slice(); + + let expression_s_16 = &mut expression.s[16]; + let expression_c_15 = &mut expression_c[15]; + let mut expression_bit_16 = expression_s_16.clone(); + expression_bit_16 ^= expression_c_15.as_slice(); + expression_bit_16 ^= expression_carry.as_slice(); + *expression_s_16 ^= expression_carry.as_slice(); + *expression_c_15 ^= expression_carry.as_slice(); + + let mut final_round = and_many_grouped( + session, + vec![ + (expression_s_16.as_slice(), expression_c_15.as_slice()), + (mask_c_msb.as_slice(), mask_carry.as_slice()), + (expression_bit_16.as_slice(), mask_carry_16.as_slice()), + ], + ) + .await?; + expression_carry ^= final_round.remove(0); + let mask_carry_17 = final_round.remove(0); + let carry_into_sign = final_round.remove(0); + + let mut result = expression + .s + .pop() + .ok_or_else(|| eyre!("missing expression sign bit"))?; + result ^= expression_c[16].as_slice(); + result ^= expression_carry; + result ^= mask_carry_17.as_slice(); + result ^= carry_into_sign.as_slice(); + Ok(result) +} + +/// Project a packed replicated share onto a subset of its three underlying +/// additive components. Component bit 0 is x1, bit 1 is x2, and bit 2 is x3. +/// +/// A party with role `i` holds `(x_i, x_(i-1))` (with zero-based component +/// indices), so projection only needs to retain or clear each local half. +fn project_packed_components( + session: &Session, + packed: &[VecShare], + component_mask: u8, +) -> Result>> { + let role = session.own_role().index(); + if role >= 3 { + bail!("cannot project Rep3 components for role {role}"); + } + let previous_role = (role + 2) % 3; + let keep_a = component_mask & (1 << role) != 0; + let keep_b = component_mask & (1 << previous_role) != 0; + + Ok(packed + .iter() + .map(|plane| { + VecShare::new_vec( + plane + .iter() + .map(|share| { + Share::new( + if keep_a { share.a } else { RingElement::zero() }, + if keep_b { share.b } else { RingElement::zero() }, + ) + }) + .collect(), + ) + }) + .collect()) +} + +/// Run the anonymous-threshold adder from packed replicated components, +/// without first expanding each packed word into three mostly-zero shares. +async fn binary_add_anon_threshold_from_packed_rep3( + session: &mut Session, + mask: Vec>, + expression: Vec>, +) -> Result, Error> { + // For x = x1+x2+x3, the bitwise full-adder reduction uses + // s=x1^x2^x3 and c=(x1^x3)&(x2^x3)^x3. The packed Rep3 share already + // represents s, while the three projected operands below can be formed + // directly from its two locally-held components. + let mask_x1x3 = project_packed_components(session, &mask, 0b101)?; + let mask_x2x3 = project_packed_components(session, &mask, 0b110)?; + let mask_x3 = project_packed_components(session, &mask, 0b100)?; + + let expression_x1x3 = project_packed_components(session, &expression, 0b101)?; + let expression_x2x3 = project_packed_components(session, &expression, 0b110)?; + let expression_x3 = project_packed_components(session, &expression, 0b100)?; + + let mask = ThreeWayAdderParts { + s: mask, + x1x3: mask_x1x3, + x2x3: mask_x2x3, + x3: mask_x3, + }; + let expression = ThreeWayAdderParts { + s: expression, + x1x3: expression_x1x3, + x2x3: expression_x2x3, + x3: expression_x3, + }; + binary_add_anon_threshold_from_parts(session, mask, expression).await +} + /// Conducts a 3 party protocol to inject bits into shares of type T. /// The protocol is given in , see Section 6.2 and Protocol 22. /// @@ -1261,6 +1547,237 @@ where Ok(res) } +fn split_packed_additive_components( + session: &Session, + packed: Vec>, +) -> Result { + let packed_len = packed.len(); + let mut x1 = Vec::with_capacity(packed_len); + let mut x2 = Vec::with_capacity(packed_len); + let mut x3 = Vec::with_capacity(packed_len); + for bit_slice in packed { + let words = bit_slice.len(); + let mut x1_slice = VecShare::with_capacity(words); + let mut x2_slice = VecShare::with_capacity(words); + let mut x3_slice = VecShare::with_capacity(words); + for word in bit_slice { + let (a, b, c) = a2b_pre(session, word)?; + x1_slice.push(a); + x2_slice.push(b); + x3_slice.push(c); + } + x1.push(x1_slice); + x2.push(x2_slice); + x3.push(x3_slice); + } + Ok((x1, x2, x3)) +} + +#[inline] +fn transpose_u16_block(input: &[u16; 64]) -> [u64; 16] { + let mut result = [0_u64; 16]; + for (bit, output) in result.iter_mut().enumerate() { + *output = u64::from(input[bit]) + | (u64::from(input[16 + bit]) << 16) + | (u64::from(input[32 + bit]) << 32) + | (u64::from(input[48 + bit]) << 48); + } + + let mut mask = 0x00ff00ff00ff00ff_u64; + let mut shift = 8_u32; + while shift != 0 { + let mut index = 0; + while index < 16 { + let swap = ((result[index] >> shift) ^ result[index + shift as usize]) & mask; + result[index + shift as usize] ^= swap; + result[index] ^= swap << shift; + index = (index + shift as usize + 1) & !(shift as usize); + } + shift >>= 1; + mask ^= mask << shift; + } + result +} + +#[inline] +fn transpose_u32_block(input: &[u32; 64]) -> [u64; 32] { + let mut result = [0_u64; 32]; + for (bit, output) in result.iter_mut().enumerate() { + *output = u64::from(input[bit]) | (u64::from(input[32 + bit]) << 32); + } + + let mut mask = 0x0000ffff0000ffff_u64; + let mut shift = 16_u32; + while shift != 0 { + let mut index = 0; + while index < 32 { + let swap = ((result[index] >> shift) ^ result[index + shift as usize]) & mask; + result[index + shift as usize] ^= swap; + result[index] ^= swap << shift; + index = (index + shift as usize + 1) & !(shift as usize); + } + shift >>= 1; + mask ^= mask << shift; + } + result +} + +/// Bit-slice refreshed, interleaved `(code, mask)` components directly into +/// packed Rep3 words. This avoids allocating dense scalar Rep3 shares and +/// transposing their two halves together only to split them again. +fn pack_anon_stats_components( + local: &[RingElement], + previous: &[RingElement], +) -> Result { + if local.len() != previous.len() { + bail!("local and previous component batches must have equal lengths"); + } + if !local.len().is_multiple_of(2) { + bail!("anonymous-threshold input must contain interleaved code/mask pairs"); + } + + const EXPRESSION_BITS: usize = 18; + const EXPRESSION_MASK: u32 = (1_u32 << EXPRESSION_BITS) - 1; + let comparisons = local.len() / 2; + let packed_words = comparisons.div_ceil(64); + let mut masks = (0..16) + .map(|_| VecShare::with_capacity(packed_words)) + .collect::>(); + let mut expressions = (0..EXPRESSION_BITS) + .map(|_| VecShare::with_capacity(packed_words)) + .collect::>(); + + for word in 0..packed_words { + let mut local_masks = [0_u16; 64]; + let mut previous_masks = [0_u16; 64]; + let mut local_expressions = [0_u32; 64]; + let mut previous_expressions = [0_u32; 64]; + let start = word * 64; + let end = (start + 64).min(comparisons); + + for (lane, comparison) in (start..end).enumerate() { + let index = comparison * 2; + let local_code = u32::from(local[index].0); + let local_mask = local[index + 1].0; + let previous_code = u32::from(previous[index].0); + let previous_mask = previous[index + 1].0; + + local_masks[lane] = local_mask; + previous_masks[lane] = previous_mask; + local_expressions[lane] = + ((local_code << 2).wrapping_sub(u32::from(local_mask))) & EXPRESSION_MASK; + previous_expressions[lane] = + ((previous_code << 2).wrapping_sub(u32::from(previous_mask))) & EXPRESSION_MASK; + } + + let local_masks = transpose_u16_block(&local_masks); + let previous_masks = transpose_u16_block(&previous_masks); + for (plane, (a, b)) in masks + .iter_mut() + .zip(local_masks.into_iter().zip(previous_masks)) + { + plane.push(Share::new(RingElement(a), RingElement(b))); + } + + let local_expressions = transpose_u32_block(&local_expressions); + let previous_expressions = transpose_u32_block(&previous_expressions); + for (plane, (a, b)) in expressions.iter_mut().zip( + local_expressions + .into_iter() + .zip(previous_expressions) + .take(EXPRESSION_BITS), + ) { + plane.push(Share::new(RingElement(a), RingElement(b))); + } + } + + Ok((masks, expressions)) +} + +/// Extract the fixed anonymous-threshold result directly from refreshed local +/// and previous additive components in interleaved `(code, mask)` order. +pub(crate) async fn extract_anon_stats_msb_batch_from_components( + session: &mut Session, + local: &[RingElement], + previous: &[RingElement], +) -> Result>> { + if local.is_empty() && previous.is_empty() { + return Ok(Vec::new()); + } + let result_len = local.len() / 2; + let (masks, expressions) = pack_anon_stats_components(local, previous)?; + let bit_17 = binary_add_anon_threshold_from_packed_rep3(session, masks, expressions).await?; + let mut result = bit_17.convert_to_bits(); + result.truncate(result_len); + Ok(result.inner()) +} + +/// Extract the sign of the anonymous-statistics FHD threshold expression +/// directly from replicated 16-bit code and mask dot products. +/// +/// For the fixed anonymous threshold `t = 0.375`, `A = 2^14` and `B = 2^16`: +/// +/// `code * B - mask * A = (4 * code - mask) * 2^14`. +/// +/// Its 32-bit sign is therefore bit 17 of the expression in the 18-bit ring. +/// The mask is shared modulo 2^16, so the two overflow bits from adding its +/// three additive components are carried directly into bits 16 and 17. This +/// avoids lifting every mask to arithmetic u32 shares and avoids bit-injecting +/// the lift corrections. The resulting bit is exactly the same one returned +/// by `extract_msb_batch` after the generic lift/multiply/subtract path. +#[instrument(level = "trace", target = "searcher::network", skip_all)] +pub async fn extract_anon_stats_msb_batch( + session: &mut Session, + code_dots: &[Share], + mask_dots: &[Share], +) -> Result>> { + if code_dots.len() != mask_dots.len() { + bail!("code and mask batches must have equal lengths"); + } + if code_dots.is_empty() { + return Ok(Vec::new()); + } + const BITS: usize = 18; + const RING_MASK: u32 = (1_u32 << BITS) - 1; + let result_len = code_dots.len(); + + // Compute the quotient q in x0+x1+x2 = mask + q*2^16 while evaluating the + // 18-bit threshold expression. Equal-depth gates from the two independent + // adders are combined into the same communication rounds. + let packed_masks = VecShare::new_vec(mask_dots.to_vec()).transpose_pack_u64(); + let (mask_x1, mask_x2, mask_x3) = split_packed_additive_components(session, packed_masks)?; + + // Add the three local components of 4*code-mask modulo 2^18. Reducing + // each component independently is valid in the target power-of-two ring. + let raw_expression = code_dots + .iter() + .zip(mask_dots) + .map(|(code, mask)| { + let component = |code: RingElement, mask: RingElement| { + RingElement(((u32::from(code.0) << 2).wrapping_sub(u32::from(mask.0))) & RING_MASK) + }; + Share::new(component(code.a, mask.a), component(code.b, mask.b)) + }) + .collect::>(); + let mut packed_expression = VecShare::new_vec(raw_expression).transpose_pack_u64(); + packed_expression.truncate(BITS); + let (expression_x1, expression_x2, expression_x3) = + split_packed_additive_components(session, packed_expression)?; + let bit_17 = binary_add_anon_threshold_fused( + session, + mask_x1, + mask_x2, + mask_x3, + expression_x1, + expression_x2, + expression_x3, + ) + .await?; + let mut result = bit_17.convert_to_bits(); + result.truncate(result_len); + Ok(result.inner()) +} + /// Opens a vector of binary additive replicated secret shares as described in the ABY3 framework. /// /// In particular, each party holds a share of the form `(a, b)` where `a` and `b` are already known to the next and previous parties, respectively. diff --git a/ampc-actor-utils/src/protocol/fhd_ops.rs b/ampc-actor-utils/src/protocol/fhd_ops.rs index c580d789..4c2257b2 100644 --- a/ampc-actor-utils/src/protocol/fhd_ops.rs +++ b/ampc-actor-utils/src/protocol/fhd_ops.rs @@ -1,6 +1,6 @@ use ampc_secret_sharing::{ shares::{bit::Bit, DistanceShare, VecShare}, - Share, + RingElement, Share, }; use eyre::Result; use tracing::instrument; @@ -8,11 +8,61 @@ use tracing::instrument; use crate::{ execution::session::Session, protocol::{ - binary::{bit_inject, extract_msb_batch, open_bin}, - ops::{conditionally_select_distance, reshare_products, DistancePair, B}, + binary::{ + bit_inject, extract_anon_stats_msb_batch, extract_anon_stats_msb_batch_from_components, + extract_msb_batch, lift, mul_lift_2k_to_32, open_bin, + }, + ops::{ + conditionally_select_distance, galois_ring_to_rep3_components, reshare_products, + DistancePair, B, + }, }, }; +pub type FhdDotSharePair = (Vec>, Vec>); + +/// Refreshed Rep3 dot products retained by the fused exact-scan threshold +/// path. Scalar [`Share`] values are materialized only for public candidate +/// indices after the dense anonymous-statistics comparison has completed. +#[derive(Debug)] +pub struct FusedFhdDotShares { + local: Vec>, + previous: Vec>, +} + +impl FusedFhdDotShares { + /// Number of interleaved `(code, mask)` comparisons retained. + pub fn len(&self) -> usize { + self.local.len() / 2 + } + + pub fn is_empty(&self) -> bool { + self.local.is_empty() + } + + /// Reconstruct selected scalar Rep3 code and mask shares. Indices address + /// comparisons, and output order (including duplicate indices) is + /// preserved. + pub fn select(&self, indices: &[usize]) -> Result { + let mut codes = Vec::with_capacity(indices.len()); + let mut masks = Vec::with_capacity(indices.len()); + for &index in indices { + eyre::ensure!( + index < self.len(), + "dot-product index {index} is out of bounds for {} comparisons", + self.len() + ); + let offset = index * 2; + codes.push(Share::new(self.local[offset], self.previous[offset])); + masks.push(Share::new( + self.local[offset + 1], + self.previous[offset + 1], + )); + } + Ok((codes, masks)) + } +} + /// Computes the `A` term of the threshold comparison based on the formula `A = ((1. - 2. * t) * B)`. #[inline] pub fn translate_threshold_a(t: f64) -> u32 { @@ -50,6 +100,83 @@ pub async fn fhd_greater_than_threshold( extract_msb_batch(session, &diffs).await } +/// Lift only the mask-dot shares needed by the direct FHD threshold protocol. +/// +/// Unlike the generic distance path, exact linear scan never needs a lifted +/// distance or an oblivious minimum. The code dot is multiplied by `2^16`, so +/// it can be lifted locally by shifting; only the mask dot requires the MPC +/// carry-correcting lift. +#[instrument(level = "trace", target = "searcher::network", skip_all)] +pub async fn lift_fhd_mask_dots( + session: &mut Session, + mask_dots: &[Share], +) -> Result>> { + Ok(lift(session, VecShare::new_vec(mask_dots.to_vec())) + .await? + .inner()) +} + +/// Compare raw replicated `u16` dot products to an FHD threshold. +/// +/// This is the CPU counterpart of the GPU threshold-ring `lift_mul_sub` +/// protocol. All inputs are processed in one batch. `mask_dots` must have +/// already been lifted with [`lift_fhd_mask_dots`]. +#[instrument(level = "trace", target = "searcher::network", skip_all)] +pub async fn fhd_greater_than_threshold_pre_lifted_masks( + session: &mut Session, + code_dots: &[Share], + mask_dots: &[Share], + threshold_ratio: f64, +) -> Result>> { + eyre::ensure!( + code_dots.len() == mask_dots.len(), + "code and mask dot batches must have equal lengths" + ); + let a = translate_threshold_a(threshold_ratio); + let diffs = code_dots + .iter() + .zip(mask_dots) + .map(|(code_dot, mask_dot)| mul_lift_2k_to_32::<16>(code_dot) - *mask_dot * a) + .collect::>(); + extract_msb_batch(session, &diffs).await +} + +/// Refresh local interleaved Galois-ring dot contributions into Rep3 and run +/// the fixed anonymous-statistics threshold without materializing the dense +/// scalar Rep3 batch. +/// +/// The refresh is exactly the one used by `galois_ring_to_rep3`: it consumes +/// the same PRF values, sends one `VecRing16` to the next party, and receives +/// the previous party's refreshed components. The returned holder can later +/// materialize only the publicly selected candidate code and mask shares. +#[instrument(level = "trace", target = "searcher::network", skip_all)] +pub async fn fhd_greater_than_anon_stats_from_galois( + session: &mut Session, + interleaved_dots: Vec>, +) -> Result<(Vec>, FusedFhdDotShares)> { + eyre::ensure!( + interleaved_dots.len().is_multiple_of(2), + "anonymous-threshold input must contain interleaved code/mask pairs" + ); + let (local, previous) = galois_ring_to_rep3_components(session, interleaved_dots).await?; + let bits = extract_anon_stats_msb_batch_from_components(session, &local, &previous).await?; + Ok((bits, FusedFhdDotShares { local, previous })) +} + +/// Dense exact-scan comparison for the fixed anonymous-statistics threshold. +/// +/// This uses the 18-bit direct circuit rather than lifting every mask into the +/// 32-bit arithmetic ring. Strict thresholds remain on the generic path since +/// their public multiplier is not a power of two. +#[instrument(level = "trace", target = "searcher::network", skip_all)] +pub async fn fhd_greater_than_anon_stats_threshold( + session: &mut Session, + code_dots: &[Share], + mask_dots: &[Share], +) -> Result>> { + extract_anon_stats_msb_batch(session, code_dots, mask_dots).await +} + /// Computes the cross product of distances shares represented as a fraction (code_dist, mask_dist). /// The cross product is computed as (d2.code_dist * d1.mask_dist - d1.code_dist * d2.mask_dist) and the result is shared. /// @@ -141,7 +268,10 @@ mod tests { local::{generate_local_identities, LocalRuntime}, session::SessionHandles, }, - protocol::{ops::batch_signed_lift_vec, test_utils::create_array_sharing}, + protocol::{ + ops::{batch_signed_lift_vec, galois_ring_to_rep3, open_ring}, + test_utils::create_array_sharing, + }, }; use super::*; @@ -149,7 +279,7 @@ mod tests { use aes_prng::AesRng; use ampc_secret_sharing::RingElement; use eyre::{bail, Result}; - use rand::SeedableRng; + use rand::{Rng, SeedableRng}; use std::{collections::HashMap, sync::Arc}; use tokio::{sync::Mutex, task::JoinSet}; use tracing::instrument; @@ -292,7 +422,9 @@ mod tests { let n = test_cases.len(); jobs.spawn(async move { let mut session = session.lock().await; - let lifted = batch_signed_lift_vec(&mut session, shares_i).await.unwrap(); + let lifted = batch_signed_lift_vec(&mut session, shares_i.clone()) + .await + .unwrap(); let distances: Vec> = (0..n) .map(|j| DistanceShare::new(lifted[2 * j], lifted[2 * j + 1])) .collect(); @@ -300,19 +432,56 @@ mod tests { fhd_greater_than_threshold(&mut session, &distances, MATCH_THRESHOLD_RATIO) .await .unwrap(); - let opened = open_bin(&mut session, &bits).await.unwrap(); - opened + let generic = open_bin(&mut session, &bits) + .await + .unwrap() .into_iter() .map(|x| x.convert()) - .collect::>() + .collect::>(); + + let code_dots = shares_i.iter().step_by(2).copied().collect::>(); + let mask_dots = shares_i + .iter() + .skip(1) + .step_by(2) + .copied() + .collect::>(); + let lifted_masks = lift_fhd_mask_dots(&mut session, &mask_dots).await.unwrap(); + let direct_bits = fhd_greater_than_threshold_pre_lifted_masks( + &mut session, + &code_dots, + &lifted_masks, + MATCH_THRESHOLD_RATIO, + ) + .await + .unwrap(); + let direct = open_bin(&mut session, &direct_bits) + .await + .unwrap() + .into_iter() + .map(|x| x.convert()) + .collect::>(); + let anon_bits = + fhd_greater_than_anon_stats_threshold(&mut session, &code_dots, &mask_dots) + .await + .unwrap(); + let anon_direct = open_bin(&mut session, &anon_bits) + .await + .unwrap() + .into_iter() + .map(|x| x.convert()) + .collect::>(); + (generic, direct, anon_direct) }); } - let results: Vec> = jobs.join_all().await; + let results: Vec<(Vec, Vec, Vec)> = jobs.join_all().await; // All parties should agree assert_eq!(results[0], results[1]); assert_eq!(results[1], results[2]); + assert_eq!(results[0].0, results[0].1); + assert_eq!(results[0].0, results[0].2); // Check against plaintext reference for (i, (cd, md, expected)) in test_cases.into_iter().enumerate() { @@ -323,15 +492,135 @@ mod tests { }; let reference = reference_fhd_greater_than_threshold(ref_cd, md as i64); assert_eq!( - results[0][i], reference, + results[0].0[i], reference, "Reference FHD threshold mismatch for (cd={}, md={}): got {}, expected {}", - cd, md, results[0][i], reference + cd, md, results[0].0[i], reference ); assert_eq!( - results[0][i], expected, + results[0].0[i], expected, "FHD threshold mismatch for (cd={}, md={}): got {}, expected {}", - cd, md, results[0][i], expected + cd, md, results[0].0[i], expected ); } } + + #[tokio::test] + async fn direct_anon_threshold_matches_lifted_reference_randomized() { + const N: usize = 4_097; + let mut rng = AesRng::seed_from_u64(0x18_375_u64); + let flat_values = (0..2 * N).map(|_| rng.gen::()).collect::>(); + let flat_shares = create_array_sharing(&mut rng, &flat_values); + let sessions = LocalRuntime::mock_sessions_with_channel().await.unwrap(); + let mut jobs = JoinSet::new(); + + for (party, session) in sessions.into_iter().enumerate() { + let session = session.clone(); + let shares = flat_shares.of_party(party).clone(); + jobs.spawn(async move { + let mut session = session.lock().await; + let codes = shares.iter().step_by(2).copied().collect::>(); + let masks = shares + .iter() + .skip(1) + .step_by(2) + .copied() + .collect::>(); + + let lifted_masks = lift_fhd_mask_dots(&mut session, &masks).await.unwrap(); + let reference = fhd_greater_than_threshold_pre_lifted_masks( + &mut session, + &codes, + &lifted_masks, + 0.375, + ) + .await + .unwrap(); + let reference = open_bin(&mut session, &reference).await.unwrap(); + + let direct = fhd_greater_than_anon_stats_threshold(&mut session, &codes, &masks) + .await + .unwrap(); + let direct = open_bin(&mut session, &direct).await.unwrap(); + (reference, direct) + }); + } + + let results = jobs.join_all().await; + assert_eq!(results[0], results[1]); + assert_eq!(results[1], results[2]); + assert_eq!(results[0].0, results[0].1); + } + + #[tokio::test] + async fn fused_galois_anon_threshold_matches_dense_rep3_and_selects_candidates() { + const N: usize = 257; + let mut rng = AesRng::seed_from_u64(0xf053_d375_u64); + let mut values = vec![ + 125_u16, + 500_u16, // exact 0.375 threshold + 124, + 500, // immediately above threshold + u16::MAX, + 200, // negative signed code dot + 10, + 0, // zero mask + ]; + values.extend((values.len()..2 * N).map(|_| rng.gen::())); + let additive_shares = create_array_sharing(&mut rng, &values); + let selected_indices = vec![0, 1, 63, 64, N - 1, 64]; + let expected_selected_codes = selected_indices + .iter() + .map(|&index| values[index * 2]) + .collect::>(); + let expected_selected_masks = selected_indices + .iter() + .map(|&index| values[index * 2 + 1]) + .collect::>(); + + let sessions = LocalRuntime::mock_sessions_with_channel().await.unwrap(); + let mut jobs = JoinSet::new(); + for (party, session) in sessions.into_iter().enumerate() { + let session = session.clone(); + let local_dots = additive_shares + .of_party(party) + .iter() + .map(|share| share.a) + .collect::>(); + let selected_indices = selected_indices.clone(); + jobs.spawn(async move { + let mut session = session.lock().await; + + let dense = galois_ring_to_rep3(&mut session, local_dots.clone()) + .await + .unwrap(); + let dense_codes = dense.iter().step_by(2).copied().collect::>(); + let dense_masks = dense.iter().skip(1).step_by(2).copied().collect::>(); + let dense_bits = + fhd_greater_than_anon_stats_threshold(&mut session, &dense_codes, &dense_masks) + .await + .unwrap(); + let dense_open = open_bin(&mut session, &dense_bits).await.unwrap(); + + let (fused_bits, retained) = + fhd_greater_than_anon_stats_from_galois(&mut session, local_dots) + .await + .unwrap(); + assert_eq!(retained.len(), N); + assert!(!retained.is_empty()); + let fused_open = open_bin(&mut session, &fused_bits).await.unwrap(); + + let (selected_codes, selected_masks) = retained.select(&selected_indices).unwrap(); + let selected_codes = open_ring(&mut session, &selected_codes).await.unwrap(); + let selected_masks = open_ring(&mut session, &selected_masks).await.unwrap(); + (dense_open, fused_open, selected_codes, selected_masks) + }); + } + + let results = jobs.join_all().await; + assert_eq!(results[0], results[1]); + assert_eq!(results[1], results[2]); + assert_eq!(results[0].0, results[0].1); + assert_eq!(results[0].2, expected_selected_codes); + assert_eq!(results[0].3, expected_selected_masks); + } } diff --git a/ampc-actor-utils/src/protocol/ops.rs b/ampc-actor-utils/src/protocol/ops.rs index e61407c5..30e1d7c8 100644 --- a/ampc-actor-utils/src/protocol/ops.rs +++ b/ampc-actor-utils/src/protocol/ops.rs @@ -18,6 +18,7 @@ use tracing::instrument; pub type DistancePair = (DistanceShare, DistanceShare); pub type IdDistance = (Share, DistanceShare); +pub(crate) type Rep3Components16 = (Vec>, Vec>); pub const B_BITS: u64 = 16; pub const B: u32 = 1 << B_BITS; @@ -117,10 +118,10 @@ pub async fn setup_shared_seed(session: &mut NetworkSession, my_seed: PrfSeed) - /// Convert Galois Ring elements to replicated secret shares (Rep3) /// This takes a vector of ring elements and converts them to replicated shares #[instrument(level = "trace", target = "searcher::network", skip_all)] -pub async fn galois_ring_to_rep3( +pub(crate) async fn galois_ring_to_rep3_components( session: &mut Session, items: Vec>, -) -> Result>> { +) -> Result { let network = &mut session.network_session; let (prf_my_values, prf_prev_values) = session.prf.gen_rands_batch(items.len()); @@ -148,12 +149,22 @@ pub async fn galois_ring_to_rep3( _ => Err(eyre!("Error in receiving in galois_ring_to_rep3 operation")), } }?; - let res: Vec> = masked_items + Ok((masked_items, shares_b)) +} + +/// Convert Galois Ring elements to replicated secret shares (Rep3) +/// This takes a vector of ring elements and converts them to replicated shares +#[instrument(level = "trace", target = "searcher::network", skip_all)] +pub async fn galois_ring_to_rep3( + session: &mut Session, + items: Vec>, +) -> Result>> { + let (shares_a, shares_b) = galois_ring_to_rep3_components(session, items).await?; + Ok(shares_a .into_iter() .zip(shares_b) .map(|(a, b)| Share::new(a, b)) - .collect(); - Ok(res) + .collect()) } /// Compares the given distances to zero and reveal the bit "less than zero". diff --git a/ampc-actor-utils/src/sync.rs b/ampc-actor-utils/src/sync.rs index fb944701..cedb2877 100644 --- a/ampc-actor-utils/src/sync.rs +++ b/ampc-actor-utils/src/sync.rs @@ -14,7 +14,7 @@ pub async fn sync_on_job_hash( hash: &[u8; JOB_HASH_LEN], ) -> eyre::Result { tracing::info!("Synchronizing on job hash: {}", hex::encode(hash)); - let local = NetworkValue::Bytes(hash.to_vec()); + let local = NetworkValue::Bytes(hash.to_vec().into()); session.network_session.send_next(local.clone()).await?; session.network_session.send_prev(local.clone()).await?; @@ -32,7 +32,7 @@ pub async fn sync_on_job_hash( let all_bytes: Vec<&[u8]> = all .iter() .map(|nv| match nv { - NetworkValue::Bytes(b) if b.len() == JOB_HASH_LEN => Ok(b.as_slice()), + NetworkValue::Bytes(b) if b.len() == JOB_HASH_LEN => Ok(&b[..]), _ => Err(eyre::eyre!( "Unexpected network value in job hash sync (expected {JOB_HASH_LEN} bytes)" )), diff --git a/ampc-actor-utils/tests/control_channel.rs b/ampc-actor-utils/tests/control_channel.rs index 1d4ec11c..b14af454 100644 --- a/ampc-actor-utils/tests/control_channel.rs +++ b/ampc-actor-utils/tests/control_channel.rs @@ -98,7 +98,7 @@ async fn test_control_channel_send_next_recv_prev() { .await .expect("control_channel() failed"); - cc.send_next(NetworkValue::Bytes(vec![party_index as u8])) + cc.send_next(NetworkValue::Bytes(vec![party_index as u8].into())) .await .expect("send_next() failed"); @@ -106,8 +106,8 @@ async fn test_control_channel_send_next_recv_prev() { let expected = ((party_index + NUM_PARTIES - 1) % NUM_PARTIES) as u8; match received { NetworkValue::Bytes(b) => assert_eq!( - b, - vec![expected], + &b[..], + &[expected], "party {party_index}: recv_prev got wrong payload" ), other => panic!("party {party_index}: unexpected variant {other:?}"), @@ -143,7 +143,7 @@ async fn test_control_channel_send_prev_recv_next() { .await .expect("control_channel() failed"); - cc.send_prev(NetworkValue::Bytes(vec![party_index as u8])) + cc.send_prev(NetworkValue::Bytes(vec![party_index as u8].into())) .await .expect("send_prev() failed"); @@ -151,8 +151,8 @@ async fn test_control_channel_send_prev_recv_next() { let expected = ((party_index + 1) % NUM_PARTIES) as u8; match received { NetworkValue::Bytes(b) => assert_eq!( - b, - vec![expected], + &b[..], + &[expected], "party {party_index}: recv_next got wrong payload" ), other => panic!("party {party_index}: unexpected variant {other:?}"), diff --git a/ampc-anon-stats/src/server/sync.rs b/ampc-anon-stats/src/server/sync.rs index 318e03a3..8bc24c4d 100644 --- a/ampc-anon-stats/src/server/sync.rs +++ b/ampc-anon-stats/src/server/sync.rs @@ -52,7 +52,7 @@ async fn broadcast_usize(session: &mut Session, value: usize) -> eyre::Result<[u } async fn broadcast_u64(session: &mut Session, value: u64) -> eyre::Result<[u64; 3]> { - let network_value = NetworkValue::Bytes(value.to_le_bytes().to_vec()); + let network_value = NetworkValue::Bytes(value.to_le_bytes().to_vec().into()); let broadcasted_values = broadcast(session, network_value).await?; let mut result = [0u64; 3]; diff --git a/deny.toml b/deny.toml index f8e7386d..814fd395 100644 --- a/deny.toml +++ b/deny.toml @@ -11,6 +11,9 @@ ignore = [ { id = "RUSTSEC-2025-0141", reason = "bincode 1.3.3 is considered complete by maintainers" }, # rand 0.8.5 unsoundness only affects custom loggers using rand::rng() during reseed { id = "RUSTSEC-2026-0097", reason = "transitive dep from sqlx/nalgebra/statrs; we don't use custom loggers calling rand::rng()" }, + # lru 0.16.3 pop() panic-safety bug; aws-sdk-s3 ^1.127 still pins lru ^0.16.3. + # S3 uses CacheKey(String), whose Drop cannot panic, and no AWS cache operation is wrapped in catch_unwind. + { id = "RUSTSEC-2026-0253", reason = "unbumpable aws-sdk-s3 transitive dep; vulnerable panic-during-key-drop path is unreachable for its CacheKey(String)" }, # quick-xml 0.26.0 XML-parse DoS (RUSTSEC-2026-0194 quadratic attribute check, RUSTSEC-2026-0195 # NsReader namespace-alloc). Both patched only in quick-xml >= 0.41.0. Bump is infeasible: quick-xml # is reached solely via pprof 0.15 -> inferno ^0.11 (flamegraph feature in ampc-server-utils), and